EnzymeAD / EnzymeAD/Enzyme-JAX

Transpose of gather

Open
#937 0 comments 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.