Traceback (most recent call last):
File "/media/adam/shared_drive/PycharmProjects/triton_test/pallas_test.py", line 67, in <module>
app.run(main)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/absl/app.py", line 308, in run
_run_main(main, args)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/absl/app.py", line 254, in _run_main
sys.exit(main(argv))
File "/media/adam/shared_drive/PycharmProjects/triton_test/pallas_test.py", line 45, in main
dx, dw, db = f_grad(x, weight, bias)
File "/media/adam/shared_drive/PycharmProjects/triton_test/pallas_test.py", line 37, in f
return layer_norm.layer_norm(x, w, b).sum()
File "/home/adam/Downloads/jax-triton/jax_triton/pallas/ops/layer_norm.py", line 100, in layer_norm_forward
out, mean, rstd = method(x, weight, bias)
File "/home/adam/Downloads/jax-triton/jax_triton/pallas/pallas_call.py", line 397, in wrapped
out_flat = pallas_call_p.bind(
jax._src.source_info_util.JaxStackTraceBeforeTransformation: RuntimeError: Cannot call @triton.jit'd outside of the scope of a kernel
The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.
--------------------
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/media/adam/shared_drive/PycharmProjects/triton_test/pallas_test.py", line 67, in <module>
app.run(main)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/absl/app.py", line 308, in run
_run_main(main, args)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/absl/app.py", line 254, in _run_main
sys.exit(main(argv))
File "/media/adam/shared_drive/PycharmProjects/triton_test/pallas_test.py", line 45, in main
dx, dw, db = f_grad(x, weight, bias)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/pjit.py", line 238, in cache_miss
outs, out_flat, out_tree, args_flat = _python_pjit_helper(
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/pjit.py", line 185, in _python_pjit_helper
out_flat = pjit_p.bind(*args_flat, **params)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/core.py", line 2592, in bind
return self.bind_with_trace(top_trace, args, params)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/core.py", line 363, in bind_with_trace
out = trace.process_primitive(self, map(trace.full_raise, args), params)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/core.py", line 817, in process_primitive
return primitive.impl(*tracers, **params)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/pjit.py", line 1229, in _pjit_call_impl
compiled = _pjit_lower(
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/pjit.py", line 1315, in _pjit_lower
return _pjit_lower_cached(jaxpr, in_shardings, out_shardings, *args, **kwargs)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/pjit.py", line 1374, in _pjit_lower_cached
return pxla.lower_sharding_computation(
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/interpreters/pxla.py", line 2551, in lower_sharding_computation
lowering_result = mlir.lower_jaxpr_to_module(
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/interpreters/mlir.py", line 742, in lower_jaxpr_to_module
lower_jaxpr_to_fun(
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/interpreters/mlir.py", line 1037, in lower_jaxpr_to_fun
out_vals, tokens_out = jaxpr_subcomp(ctx.replace(name_stack=callee_name_stack),
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/interpreters/mlir.py", line 1172, in jaxpr_subcomp
ans = rule(rule_ctx, *map(_unwrap_singleton_ir_values, in_nodes),
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/pjit.py", line 1453, in _pjit_lowering
func = mlir.lower_jaxpr_to_fun(
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/interpreters/mlir.py", line 1037, in lower_jaxpr_to_fun
out_vals, tokens_out = jaxpr_subcomp(ctx.replace(name_stack=callee_name_stack),
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/jax/_src/interpreters/mlir.py", line 1172, in jaxpr_subcomp
ans = rule(rule_ctx, *map(_unwrap_singleton_ir_values, in_nodes),
File "/home/adam/Downloads/jax-triton/jax_triton/pallas/triton_ir_lowering.py", line 850, in pallas_call_lowering
compilation_result = compile_jaxpr(jaxpr, tuple((*in_shapes, *out_shapes)),
File "/home/adam/Downloads/jax-triton/jax_triton/pallas/triton_ir_lowering.py", line 817, in compile_jaxpr
lowering_result = lower_jaxpr_to_triton_module(jaxpr, in_shapes, grid_spec, name)
File "/home/adam/Downloads/jax-triton/jax_triton/pallas/triton_ir_lowering.py", line 173, in lower_jaxpr_to_triton_module
() = lower_jaxpr_to_triton_ir(ctx, jaxpr, block_infos, *args)
File "/home/adam/Downloads/jax-triton/jax_triton/pallas/triton_ir_lowering.py", line 212, in lower_jaxpr_to_triton_ir
outvals = rule(rule_ctx, *invals, **eqn.params)
File "/home/adam/Downloads/jax-triton/jax_triton/pallas/triton_ir_lowering.py", line 646, in _reduce_lowering
return triton_op(a, axis=axis, _builder=ctx.builder)
File "/home/adam/Downloads/triton/python/triton/runtime/jit.py", line 395, in __call__
raise RuntimeError("Cannot call @triton.jit'd outside of the scope of a kernel")
jax._src.traceback_util.UnfilteredStackTrace: RuntimeError: Cannot call @triton.jit'd outside of the scope of a kernel
The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.
--------------------
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/media/adam/shared_drive/PycharmProjects/triton_test/pallas_test.py", line 67, in <module>
app.run(main)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/absl/app.py", line 308, in run
_run_main(main, args)
File "/home/adam/anaconda3/envs/triton_test/lib/python3.9/site-packages/absl/app.py", line 254, in _run_main
sys.exit(main(argv))
File "/media/adam/shared_drive/PycharmProjects/triton_test/pallas_test.py", line 45, in main
dx, dw, db = f_grad(x, weight, bias)
File "/home/adam/Downloads/jax-triton/jax_triton/pallas/triton_ir_lowering.py", line 850, in pallas_call_lowering
compilation_result = compile_jaxpr(jaxpr, tuple((*in_shapes, *out_shapes)),
File "/home/adam/Downloads/jax-triton/jax_triton/pallas/triton_ir_lowering.py", line 817, in compile_jaxpr
lowering_result = lower_jaxpr_to_triton_module(jaxpr, in_shapes, grid_spec, name)
File "/home/adam/Downloads/jax-triton/jax_triton/pallas/triton_ir_lowering.py", line 173, in lower_jaxpr_to_triton_module
() = lower_jaxpr_to_triton_ir(ctx, jaxpr, block_infos, *args)
File "/home/adam/Downloads/jax-triton/jax_triton/pallas/triton_ir_lowering.py", line 212, in lower_jaxpr_to_triton_ir
outvals = rule(rule_ctx, *invals, **eqn.params)
File "/home/adam/Downloads/jax-triton/jax_triton/pallas/triton_ir_lowering.py", line 646, in _reduce_lowering
return triton_op(a, axis=axis, _builder=ctx.builder)
File "/home/adam/Downloads/triton/python/triton/runtime/jit.py", line 395, in __call__
raise RuntimeError("Cannot call @triton.jit'd outside of the scope of a kernel")
RuntimeError: Cannot call @triton.jit'd outside of the scope of a kernel
Process finished with exit code 1