EnzymeAD / EnzymeAD/Enzyme-JAX
Scatter Optimizations with Window Dims specified
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```julia-repl
julia> using Reactant
AssertionError("Could not find registered platform with name: \"cuda\". Available platform names are: ")
julia> x = Reactant.to_rarray([3])
1-element ConcreteRArray{Int64, 1}:
3
julia> T = Any[nothing]
1-element Vector{Any}:
nothing
julia> function ip(m, T)
@allowscalar m[1] = 2
T[1] = m
return m
end
ip (generic function with 1 method)
julia> @code_hlo ip(x, T)
module {
func.func @main(%arg0: tensor<1xi64>) -> tensor<1xi64> {
%c = stablehlo.constant dense<2> : tensor<1xi64>
%c_0 = stablehlo.constant dense<0> : tensor<1x1xi64>
%0 = "stablehlo.scatter"(%arg0, %c_0, %c) <{scatter_dimension_numbers = #stablehlo.scatter}> ({
^bb0(%arg1: tensor, %arg2: tensor):
stablehlo.return %arg2 : tensor
}) : (tensor<1xi64>, tensor<1x1xi64>, tensor<1xi64>) -> tensor<1xi64>
return %0 : tensor<1xi64>
}
}
```
Contributor guide
No contributing guide indexed for this repository
Research direction
Reproduce the Julia session shown in the issue and inspect the generated StableHLO scatter, especially its inserted_window_dims and scatter_dims_to_operand_dims attributes. No source files or tests are named; done means the requested scatter optimization with window dimensions is implemented and its generated HLO behavior is verified.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100