Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implicit broadcast within divide op is blocking MNIST training #288

Open
umalesTT opened this issue Feb 28, 2025 · 0 comments
Open

Implicit broadcast within divide op is blocking MNIST training #288

umalesTT opened this issue Feb 28, 2025 · 0 comments

Comments

@umalesTT
Copy link
Contributor

Backward pass for MLP training for MNIST with cross-entropy loss fails with error: BinaryOpType cannot be mapped to BcastOpMath
This is due to metal issue tenstorrent/tt-metal#15216 and, on a more general side, issue that binary eltwise ops support broadcasting only for second input, which fails for non-commutative ops (as explained in tenstorrent/tt-metal#14852).
Concretely, in the case of this backward pass, broadcast in division is needed for the first input (from 1x1xf32 to 4x1xf32) as may be seen below.

%72 = "ttnn.div"(%71, %70) : (tensor<1x1xf32, #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<1x1x!tt.tile<32x32, f32>, #ttnn.buffer_type<dram>>, <interleaved>>>, tensor<4x1xf32, #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<1x1x!tt.tile<32x32, f32>, #ttnn.buffer_type<dram>>, <interleaved>>>) -> tensor<4x1xf32, #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<1x1x!tt.tile<32x32, f32>, #ttnn.buffer_type<dram>>, <interleaved>>>

Cases like this should be covered in https://github.com/tenstorrent/tt-xla/blob/main/tests/jax/ops/test_divide.py, as well as in tests for other binary eltwise non-commutative ops.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant