EnzymeAD / EnzymeAD/Enzyme-JAX
Transpose of gather
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```
module @reactant_oh_mul attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<4x100xf64>, %arg1: tensor<5xui32>) -> tensor<5x100xf64> {
%c = stablehlo.constant dense<1> : tensor<5xi64>
%0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<4x100xf64>) -> tensor<100x4xf64>
%1 = stablehlo.convert %arg1 : (tensor<5xui32>) -> tensor<5xi64>
%2 = stablehlo.subtract %1, %c : tensor<5xi64>
%3 = "stablehlo.gather"(%0, %2) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<100x4xf64>, tensor<5xi64>) -> tensor<100x5xf64>
%4 = stablehlo.transpose %3, dims = [1, 0] : (tensor<100x5xf64>) -> tensor<5x100xf64>
return %4 : tensor<5x100xf64>
}
}
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.