Skip to content

Commit

Permalink
#15320: sweep expand (#15343)
Browse files Browse the repository at this point in the history
### Ticket
[#15320](#15320)

### Problem description
Need to add sweep for expand

### What's changed
Added the run function to compare torch and ttnn expand

### Checklist
- [ ] Post commit CI passes:
  • Loading branch information
yugi957 authored Nov 22, 2024
1 parent 9245f27 commit 230d53b
Showing 1 changed file with 15 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -309,4 +309,18 @@ def run(
*,
device,
):
raise Exception("Expand is not supported, TODO: implement via recursive concat with itself")
torch_tensor = torch_random(expand_specs["shape"], -10, 10, dtype=torch.bfloat16)
expanded_tensor = torch_tensor.expand(expand_specs["size"])

ttnn_tensor = ttnn.from_torch(torch_tensor, device=device, layout=layout, dtype=dtype)

start_time = start_measuring_time()
expanded_ttnn_tensor = ttnn.expand(ttnn_tensor, expand_specs["size"])
e2e_perf = stop_measuring_time(start_time)

ttnn_output_tensor = ttnn.to_torch(expanded_ttnn_tensor)

result = check_with_pcc(expanded_tensor, ttnn_output_tensor, 0.999)

return [result, e2e_perf]
# raise Exception("Expand is not supported, TODO: implement via recursive concat with itself")

0 comments on commit 230d53b

Please sign in to comment.