()
| 345 | ) |
| 346 | # TODO(yongwww): fix flaky error of "invalid device ordinal" |
| 347 | def test_conv(): |
| 348 | import jax |
| 349 | import jax.random as jrandom |
| 350 | from flax import linen as nn |
| 351 | |
| 352 | conv = nn.Conv(64, (7, 7), (2, 2), padding=[(3, 3), (3, 3)], name="conv_init") |
| 353 | input_shape = (7, 7, 5, 64) |
| 354 | |
| 355 | input_np = generate_np_inputs(input_shape)[0] |
| 356 | input_jnp = np2jnp(input_np) |
| 357 | # initialize the conv |
| 358 | weights = conv.init(jrandom.PRNGKey(0), input_jnp) |
| 359 | # get jax inference output |
| 360 | jax_output = conv.apply(weights, input_jnp) |
| 361 | |
| 362 | # assemble numpy data using weights generated above |
| 363 | kernel_np = np.asarray(weights["params"]["kernel"]) |
| 364 | bias_np = np.asarray(weights["params"]["bias"]) |
| 365 | inputs_np = [bias_np, kernel_np, input_np] |
| 366 | |
| 367 | # jit and lower to StableHLO |
| 368 | apply = functools.partial(conv.apply) |
| 369 | stablehlo_module = jax.jit(apply).lower(weights, input_jnp).compiler_ir(dialect="stablehlo") |
| 370 | |
| 371 | # convert in Relax |
| 372 | ir_mod = from_stablehlo(stablehlo_module) |
| 373 | # compile and run |
| 374 | tvm_output = get_vm_res(ir_mod, inputs_np) |
| 375 | # verify accuracy |
| 376 | tvm.testing.assert_allclose(tvm_output.numpy(), jax_output, rtol=1e-5, atol=1e-5) |
| 377 | |
| 378 | |
| 379 | if __name__ == "__main__": |
nothing calls this directly
no test coverage detected
searching dependent graphs…