EnzymeAD / EnzymeAD/Enzyme-JAX

Reverse of constant -> constant

Open
#1,646 0 comments 1 reaction 1 assignee Assigned to @Copilot View on GitHub
Dominant language
MLIR
Stars
131
Forks
53
Avg merge
1d 10h
Merged PRs (30d)
193

Description

```
module @reactant_jac2 attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<4xf64> {enzymexla.memory_effects = [], tf.aliasing_output = 1 : i32}) -> (tensor<4x2xf64>, tensor<4xf64>) attributes {enzymexla.memory_effects = []} {
%cst = stablehlo.constant dense<0.000000e+00> : tensor<1x2xf64>
%cst_0 = stablehlo.constant dense<[[0.000000e+00, 1.000000e+00]]> : tensor<1x2xf64>
%cst_1 = stablehlo.constant dense<[[1.000000e+00, 0.000000e+00]]> : tensor<1x2xf64>
%cst_2 = stablehlo.constant dense<[1.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00]> : tensor<4xf64>
%cst_3 = stablehlo.constant dense<[0.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00]> : tensor<4xf64>
%cst_4 = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 0.000000e+00]> : tensor<4xf64>
%cst_5 = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 0.000000e+00, 1.000000e+00]> : tensor<4xf64>
%0 = stablehlo.reverse %cst_2, dims = [0] : tensor<4xf64>
%1 = stablehlo.reverse %arg0, dims = [0] : tensor<4xf64>
%2 = stablehlo.reverse %cst_3, dims = [0] : tensor<4xf64>
%3 = stablehlo.reverse %1, dims = [0] {enzymexla.finite = [#enzymexla]} : tensor<4xf64>
%4 = stablehlo.reverse %cst_4, dims = [0] : tensor<4xf64>
%5 = stablehlo.reverse %3, dims = [0] {enzymexla.finite = [#enzymexla]} : tensor<4xf64>
%6 = stablehlo.reverse %cst_5, dims = [0] : tensor<4xf64>
%7 = stablehlo.reverse %5, dims = [0] : tensor<4xf64>
%8 = stablehlo.slice %arg0 [0:2] : (tensor<4xf64>) -> tensor<2xf64>
%9 = stablehlo.cosine %8 : tensor<2xf64>
%10 = stablehlo.reshape %9 : (tensor<2xf64>) -> tensor<1x2xf64>
%11 = stablehlo.multiply %cst_1, %10 : tensor<1x2xf64>
%12 = stablehlo.slice %1 [0:2] : (tensor<4xf64>) -> tensor<2xf64>
%13 = stablehlo.reshape %12 : (tensor<2xf64>) -> tensor<1x2xf64>
%14 = stablehlo.cosine %13 : tensor<1x2xf64>
%15 = stablehlo.multiply %cst_0, %14 : tensor<1x2xf64>
%16 = stablehlo.slice %3 [0:2] : (tensor<4xf64>) -> tensor<2xf64>
%17 = stablehlo.reshape %16 : (tensor<2xf64>) -> tensor<1x2xf64>
%18 = stablehlo.cosine %17 {enzymexla.finite = [#enzymexla], enzymexla.no_nan = [#enzymexla]} : tensor<1x2xf64>
%19 = stablehlo.multiply %cst, %18 {enzymexla.no_nan = [#enzymexla]} : tensor<1x2xf64>
%20 = stablehlo.slice %5 [0:2] : (tensor<4xf64>) -> tensor<2xf64>
%21 = stablehlo.reshape %20 : (tensor<2xf64>) -> tensor<1x2xf64>
%22 = stablehlo.cosine %21 {enzymexla.finite = [#enzymexla], enzymexla.no_nan = [#enzymexla]} : tensor<1x2xf64>
%23 = stablehlo.multiply %cst, %22 {enzymexla.no_nan = [#enzymexla]} : tensor<1x2xf64>
%24 = stablehlo.concatenate %11, %15, %19, %23, dim = 0 : (tensor<1x2xf64>, tensor<1x2xf64>, tensor<1x2xf64>, tensor<1x2xf64>) -> tensor<4x2xf64>
%25 = stablehlo.slice %7 [0:2] : (tensor<4xf64>) -> tensor<2xf64>
%26 = stablehlo.reshape %25 : (tensor<2xf64>) -> tensor<1x2xf64>
%27 = stablehlo.concatenate %13, %17, %21, %26, dim = 0 : (tensor<1x2xf64>, tensor<1x2xf64>, tensor<1x2xf64>, tensor<1x2xf64>) -> tensor<4x2xf64>
%28 = stablehlo.sine %27 : tensor<4x2xf64>
%29 = stablehlo.slice %0 [0:2] : (tensor<4xf64>) -> tensor<2xf64>
%30 = stablehlo.reshape %29 : (tensor<2xf64>) -> tensor<1x2xf64>
%31 = stablehlo.slice %2 [0:2] : (tensor<4xf64>) -> tensor<2xf64>
%32 = stablehlo.reshape %31 : (tensor<2xf64>) -> tensor<1x2xf64>
%33 = stablehlo.slice %4 [0:2] : (tensor<4xf64>) -> tensor<2xf64>
%34 = stablehlo.reshape %33 : (tensor<2xf64>) -> tensor<1x2xf64>
%35 = stablehlo.slice %6 [0:2] : (tensor<4xf64>) -> tensor<2xf64>
%36 = stablehlo.reshape %35 : (tensor<2xf64>) -> tensor<1x2xf64>
%37 = stablehlo.concatenate %30, %32, %34, %36, dim = 0 : (tensor<1x2xf64>, tensor<1x2xf64>, tensor<1x2xf64>, tensor<1x2xf64>) -> tensor<4x2xf64>
%38 = stablehlo.multiply %28, %37 : tensor<4x2xf64>
%39 = stablehlo.subtract %24, %38 : tensor<4x2xf64>
return %39, %7 : tensor<4x2xf64>, tensor<4xf64>
}
}
```

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.