EnzymeAD / EnzymeAD/Enzyme-JAX
Simplify Reduce(Reduce(...))
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```mlir
module @reactant_fn attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<5x4x3xf64>) -> tensor<1x4x1xf64> {
%cst = stablehlo.constant dense<0.000000e+00> : tensor
%0 = stablehlo.multiply %arg0, %arg0 : tensor<5x4x3xf64>
%1 = stablehlo.reduce(%0 init: %cst) applies stablehlo.add across dimensions = [2] : (tensor<5x4x3xf64>, tensor) -> tensor<5x4xf64>
%2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<5x4xf64>) -> tensor<4x5xf64>
%3 = stablehlo.reshape %2 : (tensor<4x5xf64>) -> tensor<1x4x5xf64>
%4 = stablehlo.reduce(%3 init: %cst) applies stablehlo.add across dimensions = [2] : (tensor<1x4x5xf64>, tensor) -> tensor<1x4xf64>
%5 = stablehlo.reshape %4 : (tensor<1x4xf64>) -> tensor<1x4x1xf64>
return %5 : tensor<1x4x1xf64>
}
}
```
Original Code for reference
```julia
function fn(x)
return sum(sum(x .^ 2; dims=1);dims=3)
end
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.