EnzymeAD / EnzymeAD/Enzyme-JAX

conditionally perform convert(concat) -> concat(convert)

Open
#1,340 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

```julia
julia> function fn(x)
Float32.(vcat(x[1:4], -x[5:8]))
end
fn (generic function with 1 method)

julia> @code_hlo fn(Reactant.to_rarray(rand(BFloat16, 8)))
module @reactant_fn attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<8xbf16>) -> tensor<8xf32> {
%0 = stablehlo.slice %arg0 [0:4] : (tensor<8xbf16>) -> tensor<4xbf16>
%1 = stablehlo.slice %arg0 [4:8] : (tensor<8xbf16>) -> tensor<4xbf16>
%2 = stablehlo.negate %1 : tensor<4xbf16>
%3 = stablehlo.convert %0 : (tensor<4xbf16>) -> tensor<4xf32>
%4 = stablehlo.convert %2 : (tensor<4xbf16>) -> tensor<4xf32>
%5 = stablehlo.concatenate %3, %4, dim = 0 : (tensor<4xf32>, tensor<4xf32>) -> tensor<8xf32>
return %5 : tensor<8xf32>
}
}
```

This might be good with we are reducing precision but when increasing precision we will end up increasing memory requirements

Contributor guide

No contributing guide indexed for this repository

Research direction

Reproduce the Julia fn example using Reactant.to_rarray and inspect its @code_hlo output. Start by tracing the handling of stablehlo.convert and stablehlo.concatenate; completion should conditionally use the requested ordering while avoiding increased memory when precision is widened, with coverage for the shown case.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
compilers, performance
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
32/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.