EnzymeAD / EnzymeAD/Reactant.jl

`hlo_call` and `@code_hlo` use different array shapes

Open
#1,448 7 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
370
Forks
74
Avg merge
18h 47m
Merged PRs (30d)
30

Description

Generating HLO code then compiling fails due to array shape mismatch.
I have a similar issue for HLO code generated by JAX.
I think this was working in April or so but cannot find the precise version of Reactant I was using.

version: julia 1.10.9, Reactant 0.2.143 (Google Colab)

```
using Reactant: Ops, to_rarray, @compile
using Random

a = to_rarray(randn(Float32, 2,3))
b = copy(a)

@compile(
Ops.hlo_call(
"""
module {
func.func @main(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf32>) -> tensor<2x3xf32> {
%0 = stablehlo.add %arg0, %arg1 : tensor<2x3xf32>
return %0 : tensor<2x3xf32>
}
}
""", a, b)
) # works

f(a,b) = a.*b
@compile f(a,b) # works

code = string(@code_hlo f(a,b))
println(code)

f2(a,b) = Ops.hlo_call(code, a, b)
@compile f2(a,b) # fails
```

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.