Simplify some tests since sharding propagation is in place

[wujingyue]

No description provided.

[wujingyue]

!test

[github-actions]

Review updated until commit 21834ae

Description

• Simplify tests by removing redundant operations

• Remove unnecessary sharding for output tensors

• Update test cases to reflect changes

---

Changes walkthrough 📝

	Relevant files
Tests	test_matmul.py Simplify multidevice tests                                                              tests/python/multidevice/test_matmul.py Removed sharding and parallelization for output tensors in multidevice_schedule Updated loops to exclude output tensors from device mesh setting +2/-8

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Output Sharding The PR removes the sharding logic for the output tensor (self.out). This could lead to performance regressions or incorrect results if the output tensor was previously sharded for performance reasons. for t in [self.inp, self.weight, self.bias]: self.sched._set_device_mesh(t, mesh) # Shard N for weight (N, K) and bias (N) for t in [self.weight, self.bias]: Code Duplication The loop for setting the device mesh is duplicated in both multidevice_schedule methods. This could be refactored to avoid code duplication. for t in [self.inp, self.weight, self.bias]: self.sched._set_device_mesh(t, mesh) Tensor Definition The tensor definitions in definition methods have different shapes and biases. Ensure that these changes do not affect the intended functionality and performance of the tests. self.inp = self.define_tensor([-1, -1, d * e]) self.weight = self.define_tensor([e, d * e]) self.out = self.ops.linear(self.inp, self.weight, None) self.add_output(self.out)

[wujingyue]

!test