EnzymeAD / EnzymeAD/Enzyme-JAX

Simplify Reduce(Reduce(...))

Open
#867 1 comment 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.