How to call Python bindings?

[Axel Wong]

Hi everyone,

I would like to use Python bindings stablehlo.reduce_window() with stablehlo.maximum() to create the following IR in text. 

%9 = "stablehlo.reduce_window"(%8, %4) ({
^bb0(%arg11: tensor<f32>, %arg12: tensor<f32>):
  %30 = stablehlo.maximum %arg11, %arg12 : tensor<f32>
  stablehlo.return %30 : tensor<f32>
}) {base_dilations = dense<1> : tensor<4xi64>, padding = dense<0> : tensor<4x2xi64>, window_dilations = dense<1> : tensor<4xi64>, window_dimensions = dense<[1, 1, 2, 2]> : tensor<4xi64>, window_strides = dense<[1, 1, 2, 2]> : tensor<4xi64>} : (tensor<1x6x28x28xf32>, tensor<f32>) -> tensor<1x6x14x14xf32>

How would the Python code looks like?

I would appreciate any help and ideas.

Thanks,

Axel

[Peter Hawkins]

Hi...

Here's how JAX does it:

https://github.com/google/jax/blob/0a5b2e17c345c6a5846399a61697f38696f0646b/jax/_src/interpreters/mlir.py#L2736

Does that answer your question?

Peter

--
You received this message because you are subscribed to the Google Groups "OpenXLA Discuss" group.
To unsubscribe from this group and stop receiving emails from it, send an email to openxla-discu...@openxla.org.
To view this discussion on the web visit https://groups.google.com/a/openxla.org/d/msgid/openxla-discuss/611e1ae0-ff50-41fd-9e24-931c3ee1dd21n%40openxla.org.
For more options, visit https://groups.google.com/a/openxla.org/d/optout.

[Axel Wong]

Thanks.

I've managed to find a solution as follows.

mp1 = stablehlo.reduce_window([RankedTensorType.get((1,6,14,14), F32Type.get())], [relu1], [vc4], ir.DenseI64ArrayAttr.get([1,1,2,2]),
                                    window_strides=ir.DenseI64ArrayAttr.get([1,1,2,2]),
                                    base_dilations=ir.DenseI64ArrayAttr.get([1,1,1,1]),
                                    window_dilations=ir.DenseI64ArrayAttr.get([1,1,1,1]),
                                    padding=ir.DenseI64ArrayAttr.get([0,0,0,0]))

block0 = mp1.owner.regions[0].blocks.append(RankedTensorType.get(tuple(), F32Type.get()), RankedTensorType.get(tuple(), F32Type.get()))
with InsertionPoint(block0):
          arg0, arg1 = block0.arguments
          m = stablehlo.maximum(arg0, arg1)
          stablehlo.return_([m])

Axel

[Jacques Pienaar]

The generated MLIR python ones allow one to define handlers for attribute types so that you should be able to do 

[1,1,2,2]

Instead of

ir.DenseI64ArrayAttr.get([1,1,2,2])

Well at least where there is no ambiguity (e.g., if both i32 and i64 were supported here).

See https://github.com/llvm/llvm-project/blob/30c10fda2ba539e70bff4f05625ec6358c0f7502/mlir/python/mlir/ir.py#L246 for how defined, nothing special or built in and one can do same for ones own attribute types or override default.

-- Jacques 

To view this discussion on the web visit https://groups.google.com/a/openxla.org/d/msgid/openxla-discuss/04602ff1-8b4a-4988-b7ea-bdc3e7bae0e9n%40openxla.org.