Comparar texto

Encontre a diferença entre dois arquivos de texto

Real-time diff

Unified diff

Collapse lines

Highlight change

Syntax highlighting

Ferramentas

Diffchecker Desktop The most secure way to run Diffchecker. Get the Diffchecker Desktop app: your diffs never leave your computer!Get Desktop

training_step with and without vmap

Created 3 years agoDiff never expires

Lines
Total
Removed

Words
Total
Removed

To continue using this feature, upgrade to Diffchecker Pro View Pricing

151 lines

Lines
Total
Added

Words
Total
Added

To continue using this feature, upgrade to Diffchecker Pro View Pricing

HloModule jit_training_step.89, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias), {3}: (3, {}, may-alias), {4}: (4, {}, may-alias), {5}: (5, {}, may-alias), {6}: (6, {}, may-alias), {7}: (7, {}, may-alias), {8}: (8, {}, may-alias), {9}: (9, {}, may-alias), {10}: (10, {}, may-alias), {11}: (11, {}, may-alias), {12}: (12, {}, may-alias), {13}: (13, {}, may-alias), {14}: (14, {}, may-alias), {15}: (15, {}, may-alias), {16}: (16, {}, may-alias), {17}: (17, {}, may-alias), {18}: (18, {}, may-alias), {19}: (19, {}, may-alias), {20}: (20, {}, may-alias), {21}: (21, {}, may-alias), {22}: (22, {}, may-alias), {23}: (23, {}, may-alias), {24}: (24, {}, may-alias), {25}: (25, {}, may-alias), {26}: (26, {}, may-alias), {27}: (27, {}, may-alias), {28}: (28, {}, may-alias), {29}: (29, {}, may-alias), {30}: (30, {}, may-alias), {31}: (31, {}, may-alias), {32}: (32, {}, may-alias), {33}: (33, {}, may-alias), {34}: (34, {}, may-alias), {35}: (35, {}, may-alias), {36}: (36, {}, may-alias) }

%fused_computation (param_0.2: s32[], param_1.3: s32[], param_2.6: f32[], param_3.10: f32[]) -> pred[] {

%param_3.10 = f32[] parameter(3)

%compare.128 = pred[] compare(f32[] %param_3.10, f32[] %param_3.10), direction=NE, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/ne" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%constant_452 = s32[] constant(2143289344)

%constant_451 = f32[] constant(0)

%compare.127 = pred[] compare(f32[] %param_3.10, f32[] %constant_451), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%constant_450 = s32[] constant(0)

%bitcast-convert.19 = s32[] bitcast-convert(f32[] %param_3.10), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/bitcast_convert_type[new_dtype=int32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%select.121 = s32[] select(pred[] %compare.127, s32[] %constant_450, s32[] %bitcast-convert.19), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%select.120 = s32[] select(pred[] %compare.128, s32[] %constant_452, s32[] %select.121), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%compare.126 = pred[] compare(s32[] %select.120, s32[] %constant_450), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%constant_449 = u32[] constant(2147483647)

%bitcast-convert.18 = u32[] bitcast-convert(f32[] %param_3.10), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/bitcast_convert_type[new_dtype=uint32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%subtract.62 = u32[] subtract(u32[] %constant_449, u32[] %bitcast-convert.18), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%bitcast-convert.17 = s32[] bitcast-convert(u32[] %subtract.62), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/bitcast_convert_type[new_dtype=int32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%select.119 = s32[] select(pred[] %compare.126, s32[] %bitcast-convert.17, s32[] %select.120), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%param_2.6 = f32[] parameter(2)

%compare.125 = pred[] compare(f32[] %param_2.6, f32[] %param_2.6), direction=NE, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/ne" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%compare.124 = pred[] compare(f32[] %param_2.6, f32[] %constant_451), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%bitcast-convert.16 = s32[] bitcast-convert(f32[] %param_2.6), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/bitcast_convert_type[new_dtype=int32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%select.118 = s32[] select(pred[] %compare.124, s32[] %constant_450, s32[] %bitcast-convert.16), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%select.117 = s32[] select(pred[] %compare.125, s32[] %constant_452, s32[] %select.118), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%compare.123 = pred[] compare(s32[] %select.117, s32[] %constant_450), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%bitcast-convert.15 = u32[] bitcast-convert(f32[] %param_2.6), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/bitcast_convert_type[new_dtype=uint32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%subtract.61 = u32[] subtract(u32[] %constant_449, u32[] %bitcast-convert.15), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%bitcast-convert.14 = s32[] bitcast-convert(u32[] %subtract.61), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/bitcast_convert_type[new_dtype=int32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%select.116 = s32[] select(pred[] %compare.123, s32[] %bitcast-convert.14, s32[] %select.117), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%compare.122 = pred[] compare(s32[] %select.119, s32[] %select.116), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%compare.121 = pred[] compare(s32[] %select.116, s32[] %select.119), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}

%compare.120 = pred[] compare(pred[] %compare.122, pred[] %compare.121), direction=EQ

%param_0.2 = s32[] parameter(0)

%param_1.3 = s32[] parameter(1)

%compare.119 = pred[] compare(s32[] %param_0.2, s32[] %param_1.3), direction=LT

ROOT %select.115 = pred[] select(pred[] %compare.120, pred[] %compare.119, pred[] %compare.122)

}

%region_5.625.clone (Arg_0.0: f32[], Arg_1.0: f32[], Arg_2.0: s32[], Arg_3.0: s32[], p.2.lhs.clone: s32[], p.2.rhs.clone: s32[]) -> pred[] {

%region_5.628.clone (Arg_0.0: f32[], Arg_1.0: f32[], Arg_2.0: s32[], Arg_3.0: s32[], p.2.lhs.clone: s32[], p.2.rhs.clone: s32[]) -> pred[] {

%Arg_2.0 = s32[] parameter(2)

%Arg_3.0 = s32[] parameter(3)

%p.2.lhs.clone = s32[] parameter(4)

%p.2.rhs.clone = s32[] parameter(5)

%Arg_1.0 = f32[] parameter(1)

%Arg_0.0 = f32[] parameter(0)

ROOT %fusion = pred[] fusion(s32[] %p.2.lhs.clone, s32[] %p.2.rhs.clone, f32[] %Arg_1.0, f32[] %Arg_0.0), kind=kLoop, calls=%fused_computation

}

%region_27.1710 (Arg_0.1711: f32[], Arg_1.1712: f32[]) -> f32[] {

%region_27.1717 (Arg_0.1718: f32[], Arg_1.1719: f32[]) -> f32[] {

%Arg_0.1711 = f32[] parameter(0)

%Arg_0.1718 = f32[] parameter(0)

%Arg_1.1712 = f32[] parameter(1)

%Arg_1.1719 = f32[] parameter(1)

ROOT %add.1713 = f32[] add(f32[] %Arg_0.1711, f32[] %Arg_1.1712), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}

ROOT %add.1720 = f32[] add(f32[] %Arg_0.1718, f32[] %Arg_1.1719), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}

}

%fused_computation.1 (param_0.363: f32[]) -> (f32[], f32[]) {

%fused_computation.1 (param_0.369: f32[]) -> (f32[], f32[]) {

%param_0.363 = f32[] parameter(0)

%param_0.369 = f32[] parameter(0)

%constant_471 = f32[] constant(8.13802108e-05)

%multiply.396 = f32[] multiply(f32[] %param_0.363, f32[] %constant_471), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}

%multiply.396 = f32[] multiply(f32[] %param_0.369, f32[] %constant_471), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}

%constant_454 = f32[] constant(1e-10)

%maximum.59 = f32[] maximum(f32[] %multiply.396, f32[] %constant_454), metadata={op_name="jit(training_step)/jit(main)/max" source_file="/home/brent/tensorf-jax/tensorf/utils.py" source_line=8}

%log.3 = f32[] log(f32[] %maximum.59), metadata={op_name="jit(training_step)/jit(main)/log" source_file="/home/brent/tensorf-jax/tensorf/utils.py" source_line=17}

%constant_453 = f32[] constant(-4.34294462)

%multiply.232 = f32[] multiply(f32[] %log.3, f32[] %constant_453), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/utils.py" source_line=17}

ROOT %tuple.44 = (f32[], f32[]) tuple(f32[] %multiply.232, f32[] %multiply.396)

}

%region_10.1466 (Arg_0.1467: f32[], Arg_1.1468: f32[]) -> f32[] {

%region_10.1472 (Arg_0.1473: f32[], Arg_1.1474: f32[]) -> f32[] {

%Arg_0.1467 = f32[] parameter(0)

%Arg_0.1473 = f32[] parameter(0)

%Arg_1.1468 = f32[] parameter(1)

%Arg_1.1474 = f32[] parameter(1)

ROOT %add.1469 = f32[] add(f32[] %Arg_0.1467, f32[] %Arg_1.1468), metadata={op_name="jit(training_step)/jit(main)/reduce_sum[axes=(0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}

ROOT %add.1475 = f32[] add(f32[] %Arg_0.1473, f32[] %Arg_1.1474), metadata={op_name="jit(training_step)/jit(main)/reduce_sum[axes=(0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}

}

%input_fused_computation_reduce (param_0.364: f32[4096,3]) -> f32[] {

%input_fused_computation_reduce (param_0.370: f32[4096,3]) -> f32[] {

%param_0.364 = f32[4096,3]{1,0} parameter(0)

%param_0.370 = f32[4096,3]{1,0} parameter(0)

%multiply.233 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %param_0.364, f32[4096,3]{1,0} %param_0.364), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}

%multiply.233 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %param_0.370, f32[4096,3]{1,0} %param_0.370), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}

%bitcast.74 = f32[12288]{0} bitcast(f32[4096,3]{1,0} %multiply.233)

%bitcast.73 = f32[12288]{0} bitcast(f32[4096,3]{1,0} %multiply.233)

%constant_472 = f32[] constant(0)

ROOT %reduce.36 = f32[] reduce(f32[12288]{0} %bitcast.74, f32[] %constant_472), dimensions={0}, to_apply=%region_10.1466, metadata={op_name="jit(training_step)/jit(main)/reduce_sum[axes=(0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}

ROOT %reduce.36 = f32[] reduce(f32[12288]{0} %bitcast.73, f32[] %constant_472), dimensions={0}, to_apply=%region_10.1472, metadata={op_name="jit(training_step)/jit(main)/reduce_sum[axes=(0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}

}

%fused_computation.2 (param_0.12: u32[2], param_1.10: u32[2]) -> (u32[2], u32[2], u32[2]) {

%param_0.12 = u32[2]{0} parameter(0)

%param_1.10 = u32[2]{0} parameter(1)

%concatenate.100 = u32[4]{0} concatenate(u32[2]{0} %param_0.12, u32[2]{0} %param_1.10), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}

%bitcast.76 = u32[2,2]{1,0} bitcast(u32[4]{0} %concatenate.100), metadata={op_name="jit(training_step)/jit(main)/reshape[new_sizes=(2, 2) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}

%bitcast.75 = u32[2,2]{1,0} bitcast(u32[4]{0} %concatenate.100), metadata={op_name="jit(training_step)/jit(main)/reshape[new_sizes=(2, 2) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}

%slice.274 = u32[1,2]{1,0} slice(u32[2,2]{1,0} %bitcast.76), slice={[1:2], [0:2]}, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(1, 0) limit_indices=(2, 2) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}

%slice.274 = u32[1,2]{1,0} slice(u32[2,2]{1,0} %bitcast.75), slice={[1:2], [0:2]}, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(1, 0) limit_indices=(2, 2) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}

%bitcast.75 = u32[2]{0} bitcast(u32[1,2]{1,0} %slice.274), metadata={op_name="jit(training_step)/jit(main)/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}

%bitcast.74 = u32[2]{0} bitcast(u32[1,2]{1,0} %slice.274), metadata={op_name="jit(training_step)/jit(main)/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}

%slice.358.clone.1 = u32[1]{0} slice(u32[2]{0} %param_0.12), slice={[1:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(1,) limit_indices=(2,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}

%bitcast.135.clone.1 = u32[] bitcast(u32[1]{0} %slice.358.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}

%bitcast.133.clone.1 = u32[] bitcast(u32[1]{0} %slice.358.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}

%broadcast.954.clone.1 = u32[2]{0} broadcast(u32[] %bitcast.135.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}

%broadcast.954.clone.1 = u32[2]{0} broadcast(u32[] %bitcast.133.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}

%slice.360.clone.1 = u32[1]{0} slice(u32[2]{0} %param_0.12), slice={[0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(0,) limit_indices=(1,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}

%bitcast.136.clone.1 = u32[] bitcast(u32[1]{0} %slice.360.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}

%bitcast.134.clone.1 = u32[] bitcast(u32[1]{0} %slice.360.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}

%broadcast.955.clone.1 = u32[2]{0} broadcast(u32[] %bitcast.136.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}

%broadcast.955.clone.1 = u32[2]{0} broadcast(u32[] %bitcast.134.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}

ROOT %tuple.87 = (u32[2]{0}, u32[2]{0}, u32[2]{0}) tuple(u32[2]{0} %bitcast.75, u32[2]{0} %broadcast.954.clone.1, u32[2]{0} %broadcast.955.clone.1)

ROOT %tuple.87 = (u32[2]{0}, u32[2]{0}, u32[2]{0}) tuple(u32[2]{0} %bitcast.74, u32[2]{0} %broadcast.954.clone.1, u32[2]{0} %broadcast.955.clone.1)

}

%param_0.13 = f32[3,16,128,128]{3,2,1,0} parameter(0)

%param_5.599 = f32[3,16,128,128]{3,2,1,0} parameter(5)

%param_5.567 = f32[3,16,128,128]{3,2,1,0} parameter(5)

%param_6.528 = f32[3,16,128,128]{3,2,1,0} parameter(6)

%param_6.501 = f32[3,16,128,128]{3,2,1,0} parameter(6)

%add.491.clone.1 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %param_5.599, f32[3,16,128,128]{3,2,1,0} %param_6.528), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%add.491.clone.1 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %param_5.567, f32[3,16,128,128]{3,2,1,0} %param_6.501), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%constant_478_clone_1 = f32[] constant(0.0100000007)

%broadcast.800.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_478_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%broadcast.799.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_478_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%multiply.241.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %add.491.clone.1, f32[3,16,128,128]{3,2,1,0} %broadcast.800.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%multiply.241.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %add.491.clone.1, f32[3,16,128,128]{3,2,1,0} %broadcast.799.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%param_7.475 = f32[3,16,128,128]{3,2,1,0} parameter(7)

%param_7.433 = f32[3,16,128,128]{3,2,1,0} parameter(7)

%constant_479_clone_1 = f32[] constant(0.9)

%broadcast.798.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_479_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%multiply.240.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %param_7.475, f32[3,16,128,128]{3,2,1,0} %broadcast.798.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%multiply.240.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %param_7.433, f32[3,16,128,128]{3,2,1,0} %broadcast.798.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%add.336.clone.1 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %multiply.241.clone.1, f32[3,16,128,128]{3,2,1,0} %multiply.240.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%param_2.996 = f32[] parameter(2)

%param_2.991 = f32[] parameter(2)

%broadcast.794 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %param_2.996), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}

%broadcast.794 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %param_2.991), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}

%constant_477_clone_1 = f32[] constant(0.1)

%broadcast.797.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_477_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}

%multiply.239.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %add.491.clone.1, f32[3,16,128,128]{3,2,1,0} %broadcast.797.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}

%multiply.238.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %multiply.239.clone.1, f32[3,16,128,128]{3,2,1,0} %multiply.239.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%constant_475_clone_1 = f32[] constant(0.01)

%broadcast.796.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_475_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%multiply.237.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %multiply.238.clone.1, f32[3,16,128,128]{3,2,1,0} %broadcast.796.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%param_4.645 = f32[3,16,128,128]{3,2,1,0} parameter(4)

%param_4.627 = f32[3,16,128,128]{3,2,1,0} parameter(4)

%constant_476_clone_1 = f32[] constant(0.99)

%broadcast.795.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_476_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%multiply.236.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %param_4.645, f32[3,16,128,128]{3,2,1,0} %broadcast.795.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%multiply.236.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %param_4.627, f32[3,16,128,128]{3,2,1,0} %broadcast.795.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%add.335.clone.1 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %multiply.237.clone.1, f32[3,16,128,128]{3,2,1,0} %multiply.236.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}

%param_1.976 = f32[] parameter(1)

%param_1.967 = f32[] parameter(1)

%broadcast.793 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %param_1.976), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}

%broadcast.793 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %param_1.967), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}

%divide.24 = f32[3,16,128,128]{3,2,1,0} divide(f32[3,16,128,128]{3,2,1,0} %add.335.clone.1, f32[3,16,128,128]{3,2,1,0} %broadcast.793), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}

%sqrt.0 = f32[3,16,128,128]{3,2,1,0} sqrt(f32[3,16,128,128]{3,2,1,0} %divide.24), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}

%constant_474 = f32[] constant(1e-08)

%broadcast.792 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_474), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}

%add.334 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %sqrt.0, f32[3,16,128,128]{3,2,1,0} %broadcast.792), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}

%multiply.235 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %broadcast.794, f32[3,16,128,128]{3,2,1,0} %add.334)

%divide.23 = f32[3,16,128,128]{3,2,1,0} divide(f32[3,16,128,128]{3,2,1,0} %add.336.clone.1, f32[3,16,128,128]{3,2,1,0} %multiply.235), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}

%param_3.881 = f32[] parameter(3)

%param_3.874 = f32[] parameter(3)

%constant_473 = f32[] constant(-0.02)

%multiply.397 = f32[] multiply(f32[] %param_3.881, f32[] %constant_473), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}

%multiply.397 = f32[] multiply(f32[] %param_3.874, f32[] %constant_473), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}

%broadcast.791 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %multiply.397), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}

%multiply.234 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %divide.23, f32[3,16,128,128]{3,2,1,0} %broadcast.791), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}

%add.333 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %param_0.13, f32[3,16,128,128]{3,2,1,0} %multiply.234), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}

ROOT %tuple.46 = (f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}) tuple(f32[3,16,128,128]{3,2,1,0} %add.333, f32[3,16,128,128]{3,2,1,0} %add.335.clone.1, f32[3,16,128,128]{3,2,1,0} %add.336.clone.1)

}

%region_24.1616 (Arg_0.1617: f32[], Arg_1.1618: f32[]) -> f32[] {

%region_24.1622 (Arg_0.1623: f32[], Arg_1.1624: f32[]) -> f32[] {

%Arg_0.1617 = f32[] parameter(0)

%Arg_0.1623 = f32[] parameter(0)

%Arg_1.1618 = f32[] parameter(1)

%Arg_1.1624 = f32[] parameter(1)

ROOT %add.1619 = f32[] add(f32[] %Arg_0.1617, f32[] %Arg_1.1618), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}

ROOT %add.1625 = f32[] add(f32[] %Arg_0.1623, f32[] %Arg_1.1624), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}

}

%input_fused_computation_scatter (param_0.553: f32[3,4096,16,221], param_1.751: f32[3,4096,221], param_2.713: f32[3,4096,221]) -> f32[3,16,128,128] {

%constant_480 = f32[] constant(0)

%broadcast.962 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_480), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 16, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%broadcast.962 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_480), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/broadcast_in_dim[shape=(3, 16, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%iota.127 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(j

%iota.135 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.p

Diffs salvos

Texto original

Abrir arquivo

%fused_computation (param_0.2: s32[], param_1.3: s32[], param_2.6: f32[], param_3.10: f32[]) -> pred[] {
  %param_3.10 = f32[] parameter(3)
  %compare.128 = pred[] compare(f32[] %param_3.10, f32[] %param_3.10), direction=NE, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/ne" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_452 = s32[] constant(2143289344)
  %constant_451 = f32[] constant(0)
  %compare.127 = pred[] compare(f32[] %param_3.10, f32[] %constant_451), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_450 = s32[] constant(0)
  %bitcast-convert.19 = s32[] bitcast-convert(f32[] %param_3.10), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/bitcast_convert_type[new_dtype=int32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %select.121 = s32[] select(pred[] %compare.127, s32[] %constant_450, s32[] %bitcast-convert.19), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %select.120 = s32[] select(pred[] %compare.128, s32[] %constant_452, s32[] %select.121), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %compare.126 = pred[] compare(s32[] %select.120, s32[] %constant_450), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_449 = u32[] constant(2147483647)
  %bitcast-convert.18 = u32[] bitcast-convert(f32[] %param_3.10), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/bitcast_convert_type[new_dtype=uint32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %subtract.62 = u32[] subtract(u32[] %constant_449, u32[] %bitcast-convert.18), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %bitcast-convert.17 = s32[] bitcast-convert(u32[] %subtract.62), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/bitcast_convert_type[new_dtype=int32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %select.119 = s32[] select(pred[] %compare.126, s32[] %bitcast-convert.17, s32[] %select.120), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %param_2.6 = f32[] parameter(2)
  %compare.125 = pred[] compare(f32[] %param_2.6, f32[] %param_2.6), direction=NE, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/ne" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %compare.124 = pred[] compare(f32[] %param_2.6, f32[] %constant_451), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %bitcast-convert.16 = s32[] bitcast-convert(f32[] %param_2.6), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/bitcast_convert_type[new_dtype=int32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %select.118 = s32[] select(pred[] %compare.124, s32[] %constant_450, s32[] %bitcast-convert.16), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %select.117 = s32[] select(pred[] %compare.125, s32[] %constant_452, s32[] %select.118), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %compare.123 = pred[] compare(s32[] %select.117, s32[] %constant_450), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %bitcast-convert.15 = u32[] bitcast-convert(f32[] %param_2.6), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/bitcast_convert_type[new_dtype=uint32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %subtract.61 = u32[] subtract(u32[] %constant_449, u32[] %bitcast-convert.15), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %bitcast-convert.14 = s32[] bitcast-convert(u32[] %subtract.61), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/bitcast_convert_type[new_dtype=int32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %select.116 = s32[] select(pred[] %compare.123, s32[] %bitcast-convert.14, s32[] %select.117), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %compare.122 = pred[] compare(s32[] %select.119, s32[] %select.116), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %compare.121 = pred[] compare(s32[] %select.116, s32[] %select.119), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %compare.120 = pred[] compare(pred[] %compare.122, pred[] %compare.121), direction=EQ
  %param_0.2 = s32[] parameter(0)
  %param_1.3 = s32[] parameter(1)
  %compare.119 = pred[] compare(s32[] %param_0.2, s32[] %param_1.3), direction=LT
  ROOT %select.115 = pred[] select(pred[] %compare.120, pred[] %compare.119, pred[] %compare.122)
}

%region_5.625.clone (Arg_0.0: f32[], Arg_1.0: f32[], Arg_2.0: s32[], Arg_3.0: s32[], p.2.lhs.clone: s32[], p.2.rhs.clone: s32[]) -> pred[] {
  %Arg_2.0 = s32[] parameter(2)
  %Arg_3.0 = s32[] parameter(3)
  %p.2.lhs.clone = s32[] parameter(4)
  %p.2.rhs.clone = s32[] parameter(5)
  %Arg_1.0 = f32[] parameter(1)
  %Arg_0.0 = f32[] parameter(0)
  ROOT %fusion = pred[] fusion(s32[] %p.2.lhs.clone, s32[] %p.2.rhs.clone, f32[] %Arg_1.0, f32[] %Arg_0.0), kind=kLoop, calls=%fused_computation
}

%region_27.1710 (Arg_0.1711: f32[], Arg_1.1712: f32[]) -> f32[] {
  %Arg_0.1711 = f32[] parameter(0)
  %Arg_1.1712 = f32[] parameter(1)
  ROOT %add.1713 = f32[] add(f32[] %Arg_0.1711, f32[] %Arg_1.1712), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
}

%fused_computation.1 (param_0.363: f32[]) -> (f32[], f32[]) {
  %param_0.363 = f32[] parameter(0)
  %constant_471 = f32[] constant(8.13802108e-05)
  %multiply.396 = f32[] multiply(f32[] %param_0.363, f32[] %constant_471), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %constant_454 = f32[] constant(1e-10)
  %maximum.59 = f32[] maximum(f32[] %multiply.396, f32[] %constant_454), metadata={op_name="jit(training_step)/jit(main)/max" source_file="/home/brent/tensorf-jax/tensorf/utils.py" source_line=8}
  %log.3 = f32[] log(f32[] %maximum.59), metadata={op_name="jit(training_step)/jit(main)/log" source_file="/home/brent/tensorf-jax/tensorf/utils.py" source_line=17}
  %constant_453 = f32[] constant(-4.34294462)
  %multiply.232 = f32[] multiply(f32[] %log.3, f32[] %constant_453), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/utils.py" source_line=17}
  ROOT %tuple.44 = (f32[], f32[]) tuple(f32[] %multiply.232, f32[] %multiply.396)
}

%region_10.1466 (Arg_0.1467: f32[], Arg_1.1468: f32[]) -> f32[] {
  %Arg_0.1467 = f32[] parameter(0)
  %Arg_1.1468 = f32[] parameter(1)
  ROOT %add.1469 = f32[] add(f32[] %Arg_0.1467, f32[] %Arg_1.1468), metadata={op_name="jit(training_step)/jit(main)/reduce_sum[axes=(0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
}

%input_fused_computation_reduce (param_0.364: f32[4096,3]) -> f32[] {
  %param_0.364 = f32[4096,3]{1,0} parameter(0)
  %multiply.233 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %param_0.364, f32[4096,3]{1,0} %param_0.364), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %bitcast.74 = f32[12288]{0} bitcast(f32[4096,3]{1,0} %multiply.233)
  %constant_472 = f32[] constant(0)
  ROOT %reduce.36 = f32[] reduce(f32[12288]{0} %bitcast.74, f32[] %constant_472), dimensions={0}, to_apply=%region_10.1466, metadata={op_name="jit(training_step)/jit(main)/reduce_sum[axes=(0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
}

%fused_computation.2 (param_0.12: u32[2], param_1.10: u32[2]) -> (u32[2], u32[2], u32[2]) {
  %param_0.12 = u32[2]{0} parameter(0)
  %param_1.10 = u32[2]{0} parameter(1)
  %concatenate.100 = u32[4]{0} concatenate(u32[2]{0} %param_0.12, u32[2]{0} %param_1.10), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %bitcast.76 = u32[2,2]{1,0} bitcast(u32[4]{0} %concatenate.100), metadata={op_name="jit(training_step)/jit(main)/reshape[new_sizes=(2, 2) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %slice.274 = u32[1,2]{1,0} slice(u32[2,2]{1,0} %bitcast.76), slice={[1:2], [0:2]}, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(1, 0) limit_indices=(2, 2) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %bitcast.75 = u32[2]{0} bitcast(u32[1,2]{1,0} %slice.274), metadata={op_name="jit(training_step)/jit(main)/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %slice.358.clone.1 = u32[1]{0} slice(u32[2]{0} %param_0.12), slice={[1:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(1,) limit_indices=(2,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %bitcast.135.clone.1 = u32[] bitcast(u32[1]{0} %slice.358.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %broadcast.954.clone.1 = u32[2]{0} broadcast(u32[] %bitcast.135.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %slice.360.clone.1 = u32[1]{0} slice(u32[2]{0} %param_0.12), slice={[0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(0,) limit_indices=(1,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %bitcast.136.clone.1 = u32[] bitcast(u32[1]{0} %slice.360.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %broadcast.955.clone.1 = u32[2]{0} broadcast(u32[] %bitcast.136.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  ROOT %tuple.87 = (u32[2]{0}, u32[2]{0}, u32[2]{0}) tuple(u32[2]{0} %bitcast.75, u32[2]{0} %broadcast.954.clone.1, u32[2]{0} %broadcast.955.clone.1)
}

%fused_computation.3 (param_0.13: f32[3,16,128,128], param_1.976: f32[], param_2.996: f32[], param_3.881: f32[], param_4.645: f32[3,16,128,128], param_5.599: f32[3,16,128,128], param_6.528: f32[3,16,128,128], param_7.475: f32[3,16,128,128]) -> (f32[3,16,128,128], f32[3,16,128,128], f32[3,16,128,128]) {
  %param_0.13 = f32[3,16,128,128]{3,2,1,0} parameter(0)
  %param_5.599 = f32[3,16,128,128]{3,2,1,0} parameter(5)
  %param_6.528 = f32[3,16,128,128]{3,2,1,0} parameter(6)
  %add.491.clone.1 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %param_5.599, f32[3,16,128,128]{3,2,1,0} %param_6.528), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_478_clone_1 = f32[] constant(0.0100000007)
  %broadcast.800.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_478_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.241.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %add.491.clone.1, f32[3,16,128,128]{3,2,1,0} %broadcast.800.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_7.475 = f32[3,16,128,128]{3,2,1,0} parameter(7)
  %constant_479_clone_1 = f32[] constant(0.9)
  %broadcast.798.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_479_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.240.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %param_7.475, f32[3,16,128,128]{3,2,1,0} %broadcast.798.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.336.clone.1 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %multiply.241.clone.1, f32[3,16,128,128]{3,2,1,0} %multiply.240.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.996 = f32[] parameter(2)
  %broadcast.794 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %param_2.996), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_477_clone_1 = f32[] constant(0.1)
  %broadcast.797.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_477_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.239.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %add.491.clone.1, f32[3,16,128,128]{3,2,1,0} %broadcast.797.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.238.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %multiply.239.clone.1, f32[3,16,128,128]{3,2,1,0} %multiply.239.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_475_clone_1 = f32[] constant(0.01)
  %broadcast.796.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_475_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.237.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %multiply.238.clone.1, f32[3,16,128,128]{3,2,1,0} %broadcast.796.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.645 = f32[3,16,128,128]{3,2,1,0} parameter(4)
  %constant_476_clone_1 = f32[] constant(0.99)
  %broadcast.795.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_476_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.236.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %param_4.645, f32[3,16,128,128]{3,2,1,0} %broadcast.795.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.335.clone.1 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %multiply.237.clone.1, f32[3,16,128,128]{3,2,1,0} %multiply.236.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.976 = f32[] parameter(1)
  %broadcast.793 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %param_1.976), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.24 = f32[3,16,128,128]{3,2,1,0} divide(f32[3,16,128,128]{3,2,1,0} %add.335.clone.1, f32[3,16,128,128]{3,2,1,0} %broadcast.793), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.0 = f32[3,16,128,128]{3,2,1,0} sqrt(f32[3,16,128,128]{3,2,1,0} %divide.24), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_474 = f32[] constant(1e-08)
  %broadcast.792 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_474), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.334 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %sqrt.0, f32[3,16,128,128]{3,2,1,0} %broadcast.792), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.235 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %broadcast.794, f32[3,16,128,128]{3,2,1,0} %add.334)
  %divide.23 = f32[3,16,128,128]{3,2,1,0} divide(f32[3,16,128,128]{3,2,1,0} %add.336.clone.1, f32[3,16,128,128]{3,2,1,0} %multiply.235), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.881 = f32[] parameter(3)
  %constant_473 = f32[] constant(-0.02)
  %multiply.397 = f32[] multiply(f32[] %param_3.881, f32[] %constant_473), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.791 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %multiply.397), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.234 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %divide.23, f32[3,16,128,128]{3,2,1,0} %broadcast.791), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.333 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %param_0.13, f32[3,16,128,128]{3,2,1,0} %multiply.234), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.46 = (f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}) tuple(f32[3,16,128,128]{3,2,1,0} %add.333, f32[3,16,128,128]{3,2,1,0} %add.335.clone.1, f32[3,16,128,128]{3,2,1,0} %add.336.clone.1)
}

%region_24.1616 (Arg_0.1617: f32[], Arg_1.1618: f32[]) -> f32[] {
  %Arg_0.1617 = f32[] parameter(0)
  %Arg_1.1618 = f32[] parameter(1)
  ROOT %add.1619 = f32[] add(f32[] %Arg_0.1617, f32[] %Arg_1.1618), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter (param_0.741: f32[3,4096,221], param_1.974: f32[3,4096,221], param_2.994: f32[3,16,4096,221], param_3.879: f32[4096,221], param_4.643: f32[4096,221], param_5.596: pred[4096], param_6.521: f32[4096], param_7.466: f32[4096], param_8.381: f32[4096,221], param_9.256: f32[4096,221], param_10.177: f32[4096,221], param_11.132: f32[4096,110], param_12.111: f32[4096,110]) -> f32[3,16,128,128] {
  %constant_480 = f32[] constant(0)
  %broadcast.962 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_480), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 16, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.127 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1017 = s32[] constant(0)
  %broadcast.1443 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1017), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.974 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.143 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.974), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.90 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.143), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1015 = s32[] constant(127)
  %broadcast.1442 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1015), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.79 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1443, s32[3,4096,221]{2,1,0} %convert.90, s32[3,4096,221]{2,1,0} %broadcast.1442), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.205 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.79, s32[3,4096,221]{2,1,0} %broadcast.1443), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1012 = s32[] constant(128)
  %broadcast.1441 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1012), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.620 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.79, s32[3,4096,221]{2,1,0} %broadcast.1441), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.217 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.205, s32[3,4096,221]{2,1,0} %add.620, s32[3,4096,221]{2,1,0} %clamp.79), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.206 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.217), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.741 = f32[3,4096,221]{2,1,0} parameter(0)
  %floor.142 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_0.741), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.89 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.142), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.78 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1443, s32[3,4096,221]{2,1,0} %convert.89, s32[3,4096,221]{2,1,0} %broadcast.1442), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.204 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.78, s32[3,4096,221]{2,1,0} %broadcast.1443), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.619 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.78, s32[3,4096,221]{2,1,0} %broadcast.1441), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.216 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.204, s32[3,4096,221]{2,1,0} %add.619, s32[3,4096,221]{2,1,0} %clamp.78), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.205 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.216), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.166 = s32[3,4096,221,3]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.127, s32[3,4096,221,1]{2,1,0,3} %bitcast.206, s32[3,4096,221,1]{2,1,0,3} %bitcast.205), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.18 = s32[3,4096,221,3]{3,2,1,0} copy(s32[3,4096,221,3]{2,1,0,3} %concatenate.166), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_972 = f32[] constant(1)
  %broadcast.1419 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_972), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.172 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.974, f32[3,4096,221]{2,1,0} %floor.143), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.171 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1419, f32[3,4096,221]{2,1,0} %subtract.172), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.170 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_0.741, f32[3,4096,221]{2,1,0} %floor.142), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.169 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1419, f32[3,4096,221]{2,1,0} %subtract.170), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.439 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.171, f32[3,4096,221]{2,1,0} %subtract.169), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1418 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.439), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.994 = f32[3,16,4096,221]{3,2,1,0} parameter(2)
  %param_10.177 = f32[4096,221]{1,0} parameter(10)
  %slice.730 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_10.177), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_11.132 = f32[4096,110]{1,0} parameter(11)
  %param_12.111 = f32[4096,110]{1,0} parameter(12)
  %add.1063 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %param_11.132, f32[4096,110]{1,0} %param_12.111), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.729 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_10.177), slice={[0:4096:1], [0:219:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.45 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.729), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.1062 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %add.1063, f32[4096,110]{1,0} %reverse.45), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.263 = f32[4096,111]{1,0} concatenate(f32[4096,1]{1,0} %slice.730, f32[4096,110]{1,0} %add.1062), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.200 = f32[4096,221]{1,0} pad(f32[4096,111]{1,0} %concatenate.263, f32[] %constant_480), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.199 = f32[4096,221]{1,0} pad(f32[4096,110]{1,0} %add.1063, f32[] %constant_480), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.1061 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %pad.200, f32[4096,221]{1,0} %pad.199), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.44 = f32[4096,221]{1,0} reverse(f32[4096,221]{1,0} %add.1061), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/rev[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_8.381 = f32[4096,221]{1,0} parameter(8)
  %param_9.256 = f32[4096,221]{1,0} parameter(9)
  %slice.728 = f32[4096,220]{1,0} slice(f32[4096,221]{1,0} %param_9.256), slice={[0:4096], [0:220]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(4096, 220) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=286}
  %pad.198 = f32[4096,221]{1,0} pad(f32[4096,220]{1,0} %slice.728, f32[] %constant_972), padding=0_0x1_0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=283}
  %multiply.673 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %param_8.381, f32[4096,221]{1,0} %pad.198), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=280}
  %negate.79 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.673), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %param_3.879 = f32[4096,221]{1,0} parameter(3)
  %negate.78 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_3.879), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_5.596 = pred[4096]{0} parameter(5)
  %param_7.466 = f32[4096]{0} parameter(7)
  %param_6.521 = f32[4096]{0} parameter(6)
  %constant_1864 = f32[] constant(0.001)
  %broadcast.2334 = f32[4096]{0} broadcast(f32[] %constant_1864), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.1060 = f32[4096]{0} add(f32[4096]{0} %param_6.521, f32[4096]{0} %broadcast.2334), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.113 = f32[4096]{0} maximum(f32[4096]{0} %param_7.466, f32[4096]{0} %add.1060), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.461 = f32[4096]{0} select(pred[4096]{0} %param_5.596, f32[4096]{0} %maximum.113, f32[4096]{0} %broadcast.2334), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %broadcast.2333 = f32[4096]{0} broadcast(f32[] %constant_480), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.460 = f32[4096]{0} select(pred[4096]{0} %param_5.596, f32[4096]{0} %param_6.521, f32[4096]{0} %broadcast.2333), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.402 = f32[4096]{0} subtract(f32[4096]{0} %select.461, f32[4096]{0} %select.460), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1862 = f32[] constant(0.00452488707)
  %broadcast.2332 = f32[4096]{0} broadcast(f32[] %constant_1862), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.672 = f32[4096]{0} multiply(f32[4096]{0} %subtract.402, f32[4096]{0} %broadcast.2332), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.2331 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.672), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.671 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.78, f32[4096,221]{1,0} %broadcast.2331), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %exponential.48 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %multiply.671), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.670 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.79, f32[4096,221]{1,0} %exponential.48), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %add.1059 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %reverse.44, f32[4096,221]{1,0} %multiply.670), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.668 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %add.1059, f32[4096,221]{1,0} %broadcast.2331), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %negate.77 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.668), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_4.643 = f32[4096,221]{1,0} parameter(4)
  %constant_1861 = f32[] constant(10)
  %broadcast.2326 = f32[4096,221]{1,0} broadcast(f32[] %constant_1861), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %add.1057 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %param_4.643, f32[4096,221]{1,0} %broadcast.2326), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %constant_1860 = f32[] constant(inf)
  %broadcast.2325 = f32[4096,221]{1,0} broadcast(f32[] %constant_1860), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.359 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %add.1057, f32[4096,221]{1,0} %broadcast.2325), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2324 = f32[4096,221]{1,0} broadcast(f32[] %constant_480), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.457 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.359, f32[4096,221]{1,0} %broadcast.2324, f32[4096,221]{1,0} %add.1057), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.358 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %param_3.879, f32[4096,221]{1,0} %broadcast.2325), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.456 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.358, f32[4096,221]{1,0} %broadcast.2324, f32[4096,221]{1,0} %param_3.879), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %subtract.400 = f32[4096,221]{1,0} subtract(f32[4096,221]{1,0} %select.457, f32[4096,221]{1,0} %select.456), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %exponential.47 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %subtract.400), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %multiply.667 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.77, f32[4096,221]{1,0} %exponential.47), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(softplus)))/jit(jit_transpose(jvp(softplus)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2323 = f32[3,16,4096,221]{2,1,0,3} broadcast(f32[4096,221]{1,0} %multiply.667), dimensions={2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 16, 4096, 221) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %copy.41 = f32[3,16,4096,221]{3,2,1,0} copy(f32[3,16,4096,221]{2,1,0,3} %broadcast.2323), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 16, 4096, 221) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %multiply.666 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %param_2.994, f32[3,16,4096,221]{3,2,1,0} %copy.41), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %multiply.242 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.1418, f32[3,16,4096,221]{3,2,1,0} %multiply.666), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.42 = f32[3,16,128,128]{3,2,1,0} scatter(f32[3,16,128,128]{3,2,1,0} %broadcast.962, s32[3,4096,221,3]{3,2,1,0} %copy.18, f32[3,16,4096,221]{3,2,1,0} %multiply.242), update_window_dims={1}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_24.1616, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_23.1612 (Arg_0.1613: f32[], Arg_1.1614: f32[]) -> f32[] {
  %Arg_0.1613 = f32[] parameter(0)
  %Arg_1.1614 = f32[] parameter(1)
  ROOT %add.1615 = f32[] add(f32[] %Arg_0.1613, f32[] %Arg_1.1614), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.1 (param_0.739: f32[3,16,128,128], param_1.972: f32[3,16,128,128], param_2.992: f32[3,4096,221], param_3.873: f32[3,4096,221], param_4.635: f32[3,16,4096,221], param_5.586: f32[4096,221], param_6.509: f32[4096,221], param_7.450: pred[4096], param_8.371: f32[4096], param_9.248: f32[4096], param_10.163: f32[4096,221], param_11.118: f32[4096,221], param_12.94: f32[4096,221], param_13.44: f32[4096,110], param_14.16: f32[4096,110]) -> f32[3,16,128,128] {
  %param_0.739 = f32[3,16,128,128]{3,2,1,0} parameter(0)
  %param_1.972 = f32[3,16,128,128]{3,2,1,0} parameter(1)
  %add.337 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %param_0.739, f32[3,16,128,128]{3,2,1,0} %param_1.972), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.123 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_963 = s32[] constant(0)
  %broadcast.1411 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_963), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_3.873 = f32[3,4096,221]{2,1,0} parameter(3)
  %floor.127 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_3.873), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.82 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.127), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_961 = s32[] constant(127)
  %broadcast.1410 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_961), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.71 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1411, s32[3,4096,221]{2,1,0} %convert.82, s32[3,4096,221]{2,1,0} %broadcast.1410), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.197 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.71, s32[3,4096,221]{2,1,0} %broadcast.1411), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_960 = s32[] constant(128)
  %broadcast.1409 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_960), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.610 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.71, s32[3,4096,221]{2,1,0} %broadcast.1409), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.209 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.197, s32[3,4096,221]{2,1,0} %add.610, s32[3,4096,221]{2,1,0} %clamp.71), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.198 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.209), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.992 = f32[3,4096,221]{2,1,0} parameter(2)
  %floor.126 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_2.992), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.81 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.126), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_962 = s32[] constant(1)
  %broadcast.1407 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_962), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.608 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.81, s32[3,4096,221]{2,1,0} %broadcast.1407), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.70 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1411, s32[3,4096,221]{2,1,0} %add.608, s32[3,4096,221]{2,1,0} %broadcast.1410), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.196 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.70, s32[3,4096,221]{2,1,0} %broadcast.1411), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.607 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.70, s32[3,4096,221]{2,1,0} %broadcast.1409), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.208 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.196, s32[3,4096,221]{2,1,0} %add.607, s32[3,4096,221]{2,1,0} %clamp.70), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.197 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.208), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.162 = s32[3,4096,221,3]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.123, s32[3,4096,221,1]{3,2,1,0} %bitcast.198, s32[3,4096,221,1]{3,2,1,0} %bitcast.197), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_919 = f32[] constant(1)
  %broadcast.1365 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_919), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.156 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_3.873, f32[3,4096,221]{2,1,0} %floor.127), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.155 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1365, f32[3,4096,221]{2,1,0} %subtract.156), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.154 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_2.992, f32[3,4096,221]{2,1,0} %floor.126), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.435 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.155, f32[3,4096,221]{2,1,0} %subtract.154), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1364 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.435), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_4.635 = f32[3,16,4096,221]{3,2,1,0} parameter(4)
  %param_12.94 = f32[4096,221]{1,0} parameter(12)
  %slice.724 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_12.94), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_13.44 = f32[4096,110]{1,0} parameter(13)
  %param_14.16 = f32[4096,110]{1,0} parameter(14)
  %add.1049 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %param_13.44, f32[4096,110]{1,0} %param_14.16), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.723 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_12.94), slice={[0:4096:1], [0:219:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.41 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.723), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.1048 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %add.1049, f32[4096,110]{1,0} %reverse.41), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.261 = f32[4096,111]{1,0} concatenate(f32[4096,1]{1,0} %slice.724, f32[4096,110]{1,0} %add.1048), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_1842 = f32[] constant(0)
  %pad.194 = f32[4096,221]{1,0} pad(f32[4096,111]{1,0} %concatenate.261, f32[] %constant_1842), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.193 = f32[4096,221]{1,0} pad(f32[4096,110]{1,0} %add.1049, f32[] %constant_1842), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.1047 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %pad.194, f32[4096,221]{1,0} %pad.193), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.40 = f32[4096,221]{1,0} reverse(f32[4096,221]{1,0} %add.1047), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/rev[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_10.163 = f32[4096,221]{1,0} parameter(10)
  %param_11.118 = f32[4096,221]{1,0} parameter(11)
  %slice.722 = f32[4096,220]{1,0} slice(f32[4096,221]{1,0} %param_11.118), slice={[0:4096], [0:220]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(4096, 220) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=286}
  %pad.192 = f32[4096,221]{1,0} pad(f32[4096,220]{1,0} %slice.722, f32[] %constant_919), padding=0_0x1_0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=283}
  %multiply.657 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %param_10.163, f32[4096,221]{1,0} %pad.192), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=280}
  %negate.73 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.657), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %param_5.586 = f32[4096,221]{1,0} parameter(5)
  %negate.72 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_5.586), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_7.450 = pred[4096]{0} parameter(7)
  %param_9.248 = f32[4096]{0} parameter(9)
  %param_8.371 = f32[4096]{0} parameter(8)
  %constant_1843 = f32[] constant(0.001)
  %broadcast.2310 = f32[4096]{0} broadcast(f32[] %constant_1843), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.1046 = f32[4096]{0} add(f32[4096]{0} %param_8.371, f32[4096]{0} %broadcast.2310), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.109 = f32[4096]{0} maximum(f32[4096]{0} %param_9.248, f32[4096]{0} %add.1046), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.449 = f32[4096]{0} select(pred[4096]{0} %param_7.450, f32[4096]{0} %maximum.109, f32[4096]{0} %broadcast.2310), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %broadcast.2309 = f32[4096]{0} broadcast(f32[] %constant_1842), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.448 = f32[4096]{0} select(pred[4096]{0} %param_7.450, f32[4096]{0} %param_8.371, f32[4096]{0} %broadcast.2309), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.396 = f32[4096]{0} subtract(f32[4096]{0} %select.449, f32[4096]{0} %select.448), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1841 = f32[] constant(0.00452488707)
  %broadcast.2308 = f32[4096]{0} broadcast(f32[] %constant_1841), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.656 = f32[4096]{0} multiply(f32[4096]{0} %subtract.396, f32[4096]{0} %broadcast.2308), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.2307 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.656), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.655 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.72, f32[4096,221]{1,0} %broadcast.2307), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %exponential.44 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %multiply.655), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.654 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.73, f32[4096,221]{1,0} %exponential.44), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %add.1045 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %reverse.40, f32[4096,221]{1,0} %multiply.654), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.652 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %add.1045, f32[4096,221]{1,0} %broadcast.2307), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %negate.71 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.652), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_6.509 = f32[4096,221]{1,0} parameter(6)
  %constant_1840 = f32[] constant(10)
  %broadcast.2302 = f32[4096,221]{1,0} broadcast(f32[] %constant_1840), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %add.1043 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %param_6.509, f32[4096,221]{1,0} %broadcast.2302), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %constant_1839 = f32[] constant(inf)
  %broadcast.2301 = f32[4096,221]{1,0} broadcast(f32[] %constant_1839), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.355 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %add.1043, f32[4096,221]{1,0} %broadcast.2301), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2300 = f32[4096,221]{1,0} broadcast(f32[] %constant_1842), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.445 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.355, f32[4096,221]{1,0} %broadcast.2300, f32[4096,221]{1,0} %add.1043), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.354 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %param_5.586, f32[4096,221]{1,0} %broadcast.2301), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.444 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.354, f32[4096,221]{1,0} %broadcast.2300, f32[4096,221]{1,0} %param_5.586), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %subtract.394 = f32[4096,221]{1,0} subtract(f32[4096,221]{1,0} %select.445, f32[4096,221]{1,0} %select.444), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %exponential.43 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %subtract.394), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %multiply.651 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.71, f32[4096,221]{1,0} %exponential.43), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(softplus)))/jit(jit_transpose(jvp(softplus)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2299 = f32[3,16,4096,221]{2,1,0,3} broadcast(f32[4096,221]{1,0} %multiply.651), dimensions={2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 16, 4096, 221) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %copy.39 = f32[3,16,4096,221]{3,2,1,0} copy(f32[3,16,4096,221]{2,1,0,3} %broadcast.2299), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 16, 4096, 221) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %multiply.650 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %param_4.635, f32[3,16,4096,221]{3,2,1,0} %copy.39), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %multiply.243 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.1364, f32[3,16,4096,221]{3,2,1,0} %multiply.650), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.43 = f32[3,16,128,128]{3,2,1,0} scatter(f32[3,16,128,128]{3,2,1,0} %add.337, s32[3,4096,221,3]{3,2,1,0} %concatenate.162, f32[3,16,4096,221]{3,2,1,0} %multiply.243), update_window_dims={1}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_23.1612, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_22.1608 (Arg_0.1609: f32[], Arg_1.1610: f32[]) -> f32[] {
  %Arg_0.1609 = f32[] parameter(0)
  %Arg_1.1610 = f32[] parameter(1)
  ROOT %add.1611 = f32[] add(f32[] %Arg_0.1609, f32[] %Arg_1.1610), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.2 (param_0.737: f32[3,4096,221], param_1.970: f32[3,4096,221], param_2.990: f32[3,16,4096,221], param_3.871: f32[4096,221], param_4.633: f32[4096,221], param_5.580: pred[4096], param_6.501: f32[4096], param_7.440: f32[4096], param_8.359: f32[4096,221], param_9.232: f32[4096,221], param_10.153: f32[4096,221], param_11.111: f32[4096,110], param_12.81: f32[4096,110]) -> f32[3,16,128,128] {
  %constant_481 = f32[] constant(0)
  %broadcast.963 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_481), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 16, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.119 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_909 = s32[] constant(0)
  %broadcast.1356 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_909), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.970 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.107 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.970), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.70 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.107), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_912 = s32[] constant(1)
  %broadcast.1355 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_912), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.588 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.70, s32[3,4096,221]{2,1,0} %broadcast.1355), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_908 = s32[] constant(127)
  %broadcast.1354 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_908), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.59 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1356, s32[3,4096,221]{2,1,0} %add.588, s32[3,4096,221]{2,1,0} %broadcast.1354), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.185 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.59, s32[3,4096,221]{2,1,0} %broadcast.1356), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_907 = s32[] constant(128)
  %broadcast.1353 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_907), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.586 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.59, s32[3,4096,221]{2,1,0} %broadcast.1353), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.197 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.185, s32[3,4096,221]{2,1,0} %add.586, s32[3,4096,221]{2,1,0} %clamp.59), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.190 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.197), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.737 = f32[3,4096,221]{2,1,0} parameter(0)
  %floor.106 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_0.737), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.69 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.106), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.58 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1356, s32[3,4096,221]{2,1,0} %convert.69, s32[3,4096,221]{2,1,0} %broadcast.1354), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.184 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.58, s32[3,4096,221]{2,1,0} %broadcast.1356), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.584 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.58, s32[3,4096,221]{2,1,0} %broadcast.1353), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.196 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.184, s32[3,4096,221]{2,1,0} %add.584, s32[3,4096,221]{2,1,0} %clamp.58), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.189 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.196), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.158 = s32[3,4096,221,3]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.119, s32[3,4096,221,1]{2,1,0,3} %bitcast.190, s32[3,4096,221,1]{2,1,0,3} %bitcast.189), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.19 = s32[3,4096,221,3]{3,2,1,0} copy(s32[3,4096,221,3]{2,1,0,3} %concatenate.158), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.144 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.970, f32[3,4096,221]{2,1,0} %floor.107), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_864 = f32[] constant(1)
  %broadcast.1310 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_864), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.143 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_0.737, f32[3,4096,221]{2,1,0} %floor.106), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.142 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1310, f32[3,4096,221]{2,1,0} %subtract.143), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.431 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.144, f32[3,4096,221]{2,1,0} %subtract.142), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1308 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.431), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.990 = f32[3,16,4096,221]{3,2,1,0} parameter(2)
  %param_10.153 = f32[4096,221]{1,0} parameter(10)
  %slice.718 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_10.153), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_11.111 = f32[4096,110]{1,0} parameter(11)
  %param_12.81 = f32[4096,110]{1,0} parameter(12)
  %add.1035 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %param_11.111, f32[4096,110]{1,0} %param_12.81), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.717 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_10.153), slice={[0:4096:1], [0:219:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.37 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.717), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.1034 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %add.1035, f32[4096,110]{1,0} %reverse.37), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.259 = f32[4096,111]{1,0} concatenate(f32[4096,1]{1,0} %slice.718, f32[4096,110]{1,0} %add.1034), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.188 = f32[4096,221]{1,0} pad(f32[4096,111]{1,0} %concatenate.259, f32[] %constant_481), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.187 = f32[4096,221]{1,0} pad(f32[4096,110]{1,0} %add.1035, f32[] %constant_481), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.1033 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %pad.188, f32[4096,221]{1,0} %pad.187), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.36 = f32[4096,221]{1,0} reverse(f32[4096,221]{1,0} %add.1033), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/rev[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_8.359 = f32[4096,221]{1,0} parameter(8)
  %param_9.232 = f32[4096,221]{1,0} parameter(9)
  %slice.716 = f32[4096,220]{1,0} slice(f32[4096,221]{1,0} %param_9.232), slice={[0:4096], [0:220]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(4096, 220) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=286}
  %pad.186 = f32[4096,221]{1,0} pad(f32[4096,220]{1,0} %slice.716, f32[] %constant_864), padding=0_0x1_0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=283}
  %multiply.641 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %param_8.359, f32[4096,221]{1,0} %pad.186), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=280}
  %negate.67 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.641), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %param_3.871 = f32[4096,221]{1,0} parameter(3)
  %negate.66 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_3.871), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_5.580 = pred[4096]{0} parameter(5)
  %param_7.440 = f32[4096]{0} parameter(7)
  %param_6.501 = f32[4096]{0} parameter(6)
  %constant_1822 = f32[] constant(0.001)
  %broadcast.2286 = f32[4096]{0} broadcast(f32[] %constant_1822), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.1032 = f32[4096]{0} add(f32[4096]{0} %param_6.501, f32[4096]{0} %broadcast.2286), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.105 = f32[4096]{0} maximum(f32[4096]{0} %param_7.440, f32[4096]{0} %add.1032), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.437 = f32[4096]{0} select(pred[4096]{0} %param_5.580, f32[4096]{0} %maximum.105, f32[4096]{0} %broadcast.2286), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %broadcast.2285 = f32[4096]{0} broadcast(f32[] %constant_481), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.435 = f32[4096]{0} select(pred[4096]{0} %param_5.580, f32[4096]{0} %param_6.501, f32[4096]{0} %broadcast.2285), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.390 = f32[4096]{0} subtract(f32[4096]{0} %select.437, f32[4096]{0} %select.435), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1820 = f32[] constant(0.00452488707)
  %broadcast.2284 = f32[4096]{0} broadcast(f32[] %constant_1820), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.640 = f32[4096]{0} multiply(f32[4096]{0} %subtract.390, f32[4096]{0} %broadcast.2284), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.2283 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.640), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.639 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.66, f32[4096,221]{1,0} %broadcast.2283), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %exponential.40 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %multiply.639), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.638 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.67, f32[4096,221]{1,0} %exponential.40), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %add.1031 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %reverse.36, f32[4096,221]{1,0} %multiply.638), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.636 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %add.1031, f32[4096,221]{1,0} %broadcast.2283), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %negate.65 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.636), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_4.633 = f32[4096,221]{1,0} parameter(4)
  %constant_1819 = f32[] constant(10)
  %broadcast.2278 = f32[4096,221]{1,0} broadcast(f32[] %constant_1819), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %add.1029 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %param_4.633, f32[4096,221]{1,0} %broadcast.2278), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %constant_1818 = f32[] constant(inf)
  %broadcast.2277 = f32[4096,221]{1,0} broadcast(f32[] %constant_1818), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.351 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %add.1029, f32[4096,221]{1,0} %broadcast.2277), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2276 = f32[4096,221]{1,0} broadcast(f32[] %constant_481), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.431 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.351, f32[4096,221]{1,0} %broadcast.2276, f32[4096,221]{1,0} %add.1029), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.350 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %param_3.871, f32[4096,221]{1,0} %broadcast.2277), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.430 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.350, f32[4096,221]{1,0} %broadcast.2276, f32[4096,221]{1,0} %param_3.871), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %subtract.388 = f32[4096,221]{1,0} subtract(f32[4096,221]{1,0} %select.431, f32[4096,221]{1,0} %select.430), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %exponential.39 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %subtract.388), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %multiply.635 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.65, f32[4096,221]{1,0} %exponential.39), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(softplus)))/jit(jit_transpose(jvp(softplus)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2275 = f32[3,16,4096,221]{2,1,0,3} broadcast(f32[4096,221]{1,0} %multiply.635), dimensions={2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 16, 4096, 221) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %copy.37 = f32[3,16,4096,221]{3,2,1,0} copy(f32[3,16,4096,221]{2,1,0,3} %broadcast.2275), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 16, 4096, 221) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %multiply.634 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %param_2.990, f32[3,16,4096,221]{3,2,1,0} %copy.37), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %multiply.244 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.1308, f32[3,16,4096,221]{3,2,1,0} %multiply.634), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.44 = f32[3,16,128,128]{3,2,1,0} scatter(f32[3,16,128,128]{3,2,1,0} %broadcast.963, s32[3,4096,221,3]{3,2,1,0} %copy.19, f32[3,16,4096,221]{3,2,1,0} %multiply.244), update_window_dims={1}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_22.1608, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_21.1604 (Arg_0.1605: f32[], Arg_1.1606: f32[]) -> f32[] {
  %Arg_0.1605 = f32[] parameter(0)
  %Arg_1.1606 = f32[] parameter(1)
  ROOT %add.1607 = f32[] add(f32[] %Arg_0.1605, f32[] %Arg_1.1606), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.3 (param_0.735: f32[3,4096,221], param_1.968: f32[3,4096,221], param_2.988: f32[3,16,4096,221], param_3.865: f32[4096,221], param_4.625: f32[4096,221], param_5.570: pred[4096], param_6.489: f32[4096], param_7.424: f32[4096], param_8.349: f32[4096,221], param_9.224: f32[4096,221], param_10.139: f32[4096,221], param_11.97: f32[4096,110], param_12.64: f32[4096,110]) -> f32[3,16,128,128] {
  %constant_482 = f32[] constant(0)
  %broadcast.964 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_482), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 16, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.115 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_854 = s32[] constant(0)
  %broadcast.1299 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_854), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.968 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.87 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.968), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.58 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.87), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_853 = s32[] constant(1)
  %broadcast.1298 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_853), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.565 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.58, s32[3,4096,221]{2,1,0} %broadcast.1298), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_852 = s32[] constant(127)
  %broadcast.1297 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_852), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.47 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1299, s32[3,4096,221]{2,1,0} %add.565, s32[3,4096,221]{2,1,0} %broadcast.1297), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.173 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.47, s32[3,4096,221]{2,1,0} %broadcast.1299), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_851 = s32[] constant(128)
  %broadcast.1296 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_851), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.563 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.47, s32[3,4096,221]{2,1,0} %broadcast.1296), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.185 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.173, s32[3,4096,221]{2,1,0} %add.563, s32[3,4096,221]{2,1,0} %clamp.47), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.178 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.185), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.735 = f32[3,4096,221]{2,1,0} parameter(0)
  %floor.86 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_0.735), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.57 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.86), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.562 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.57, s32[3,4096,221]{2,1,0} %broadcast.1298), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.46 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1299, s32[3,4096,221]{2,1,0} %add.562, s32[3,4096,221]{2,1,0} %broadcast.1297), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.172 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.46, s32[3,4096,221]{2,1,0} %broadcast.1299), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.560 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.46, s32[3,4096,221]{2,1,0} %broadcast.1296), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.184 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.172, s32[3,4096,221]{2,1,0} %add.560, s32[3,4096,221]{2,1,0} %clamp.46), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.177 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.184), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.154 = s32[3,4096,221,3]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.115, s32[3,4096,221,1]{3,2,1,0} %bitcast.178, s32[3,4096,221,1]{3,2,1,0} %bitcast.177), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.132 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.968, f32[3,4096,221]{2,1,0} %floor.87), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.131 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_0.735, f32[3,4096,221]{2,1,0} %floor.86), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.426 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.132, f32[3,4096,221]{2,1,0} %subtract.131), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1223 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.426), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.988 = f32[3,16,4096,221]{3,2,1,0} parameter(2)
  %param_10.139 = f32[4096,221]{1,0} parameter(10)
  %slice.712 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_10.139), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_11.97 = f32[4096,110]{1,0} parameter(11)
  %param_12.64 = f32[4096,110]{1,0} parameter(12)
  %add.1021 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %param_11.97, f32[4096,110]{1,0} %param_12.64), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.711 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_10.139), slice={[0:4096:1], [0:219:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.33 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.711), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.1020 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %add.1021, f32[4096,110]{1,0} %reverse.33), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.257 = f32[4096,111]{1,0} concatenate(f32[4096,1]{1,0} %slice.712, f32[4096,110]{1,0} %add.1020), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.182 = f32[4096,221]{1,0} pad(f32[4096,111]{1,0} %concatenate.257, f32[] %constant_482), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.181 = f32[4096,221]{1,0} pad(f32[4096,110]{1,0} %add.1021, f32[] %constant_482), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.1019 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %pad.182, f32[4096,221]{1,0} %pad.181), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.32 = f32[4096,221]{1,0} reverse(f32[4096,221]{1,0} %add.1019), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/rev[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_8.349 = f32[4096,221]{1,0} parameter(8)
  %param_9.224 = f32[4096,221]{1,0} parameter(9)
  %slice.710 = f32[4096,220]{1,0} slice(f32[4096,221]{1,0} %param_9.224), slice={[0:4096], [0:220]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(4096, 220) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=286}
  %constant_1805 = f32[] constant(1)
  %pad.180 = f32[4096,221]{1,0} pad(f32[4096,220]{1,0} %slice.710, f32[] %constant_1805), padding=0_0x1_0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=283}
  %multiply.625 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %param_8.349, f32[4096,221]{1,0} %pad.180), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=280}
  %negate.61 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.625), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %param_3.865 = f32[4096,221]{1,0} parameter(3)
  %negate.60 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_3.865), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_5.570 = pred[4096]{0} parameter(5)
  %param_7.424 = f32[4096]{0} parameter(7)
  %param_6.489 = f32[4096]{0} parameter(6)
  %constant_1801 = f32[] constant(0.001)
  %broadcast.2262 = f32[4096]{0} broadcast(f32[] %constant_1801), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.1018 = f32[4096]{0} add(f32[4096]{0} %param_6.489, f32[4096]{0} %broadcast.2262), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.101 = f32[4096]{0} maximum(f32[4096]{0} %param_7.424, f32[4096]{0} %add.1018), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.423 = f32[4096]{0} select(pred[4096]{0} %param_5.570, f32[4096]{0} %maximum.101, f32[4096]{0} %broadcast.2262), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %broadcast.2261 = f32[4096]{0} broadcast(f32[] %constant_482), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.422 = f32[4096]{0} select(pred[4096]{0} %param_5.570, f32[4096]{0} %param_6.489, f32[4096]{0} %broadcast.2261), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.384 = f32[4096]{0} subtract(f32[4096]{0} %select.423, f32[4096]{0} %select.422), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1799 = f32[] constant(0.00452488707)
  %broadcast.2260 = f32[4096]{0} broadcast(f32[] %constant_1799), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.624 = f32[4096]{0} multiply(f32[4096]{0} %subtract.384, f32[4096]{0} %broadcast.2260), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.2259 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.624), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.623 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.60, f32[4096,221]{1,0} %broadcast.2259), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %exponential.36 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %multiply.623), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.622 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.61, f32[4096,221]{1,0} %exponential.36), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %add.1017 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %reverse.32, f32[4096,221]{1,0} %multiply.622), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.620 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %add.1017, f32[4096,221]{1,0} %broadcast.2259), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %negate.59 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.620), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_4.625 = f32[4096,221]{1,0} parameter(4)
  %constant_1798 = f32[] constant(10)
  %broadcast.2254 = f32[4096,221]{1,0} broadcast(f32[] %constant_1798), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %add.1015 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %param_4.625, f32[4096,221]{1,0} %broadcast.2254), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %constant_1797 = f32[] constant(inf)
  %broadcast.2253 = f32[4096,221]{1,0} broadcast(f32[] %constant_1797), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.347 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %add.1015, f32[4096,221]{1,0} %broadcast.2253), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2252 = f32[4096,221]{1,0} broadcast(f32[] %constant_482), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.419 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.347, f32[4096,221]{1,0} %broadcast.2252, f32[4096,221]{1,0} %add.1015), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.346 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %param_3.865, f32[4096,221]{1,0} %broadcast.2253), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.417 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.346, f32[4096,221]{1,0} %broadcast.2252, f32[4096,221]{1,0} %param_3.865), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %subtract.382 = f32[4096,221]{1,0} subtract(f32[4096,221]{1,0} %select.419, f32[4096,221]{1,0} %select.417), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %exponential.35 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %subtract.382), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %multiply.619 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.59, f32[4096,221]{1,0} %exponential.35), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(softplus)))/jit(jit_transpose(jvp(softplus)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2251 = f32[3,16,4096,221]{2,1,0,3} broadcast(f32[4096,221]{1,0} %multiply.619), dimensions={2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 16, 4096, 221) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %copy.35 = f32[3,16,4096,221]{3,2,1,0} copy(f32[3,16,4096,221]{2,1,0,3} %broadcast.2251), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 16, 4096, 221) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %multiply.618 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %param_2.988, f32[3,16,4096,221]{3,2,1,0} %copy.35), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %multiply.245 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.1223, f32[3,16,4096,221]{3,2,1,0} %multiply.618), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.45 = f32[3,16,128,128]{3,2,1,0} scatter(f32[3,16,128,128]{3,2,1,0} %broadcast.964, s32[3,4096,221,3]{3,2,1,0} %concatenate.154, f32[3,16,4096,221]{3,2,1,0} %multiply.245), update_window_dims={1}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_21.1604, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%fused_computation.7 (param_0.24: f32[3,16,128], param_1.977: f32[], param_2.998: f32[], param_3.883: f32[], param_4.647: f32[3,16,128], param_5.602: f32[3,16,128], param_6.535: f32[3,16,128], param_7.484: f32[3,16,128]) -> (f32[3,16,128], f32[3,16,128], f32[3,16,128]) {
  %param_0.24 = f32[3,16,128]{2,1,0} parameter(0)
  %param_5.602 = f32[3,16,128]{2,1,0} parameter(5)
  %param_6.535 = f32[3,16,128]{2,1,0} parameter(6)
  %add.493.clone.1 = f32[3,16,128]{2,1,0} add(f32[3,16,128]{2,1,0} %param_5.602, f32[3,16,128]{2,1,0} %param_6.535), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_488_clone_1 = f32[] constant(0.0100000007)
  %broadcast.811.clone.1 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_488_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.254.clone.1 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %add.493.clone.1, f32[3,16,128]{2,1,0} %broadcast.811.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_7.484 = f32[3,16,128]{2,1,0} parameter(7)
  %constant_489_clone_1 = f32[] constant(0.9)
  %broadcast.810.clone.1 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_489_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.253.clone.1 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %param_7.484, f32[3,16,128]{2,1,0} %broadcast.810.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.341.clone.1 = f32[3,16,128]{2,1,0} add(f32[3,16,128]{2,1,0} %multiply.254.clone.1, f32[3,16,128]{2,1,0} %multiply.253.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.998 = f32[] parameter(2)
  %broadcast.805 = f32[3,16,128]{2,1,0} broadcast(f32[] %param_2.998), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_487_clone_1 = f32[] constant(0.1)
  %broadcast.809.clone.1 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_487_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.252.clone.1 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %add.493.clone.1, f32[3,16,128]{2,1,0} %broadcast.809.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.251.clone.1 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %multiply.252.clone.1, f32[3,16,128]{2,1,0} %multiply.252.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_485_clone_1 = f32[] constant(0.01)
  %broadcast.808.clone.1 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_485_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.250.clone.1 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %multiply.251.clone.1, f32[3,16,128]{2,1,0} %broadcast.808.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.647 = f32[3,16,128]{2,1,0} parameter(4)
  %constant_486_clone_1 = f32[] constant(0.99)
  %broadcast.807.clone.1 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_486_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.249.clone.1 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %param_4.647, f32[3,16,128]{2,1,0} %broadcast.807.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.340.clone.1 = f32[3,16,128]{2,1,0} add(f32[3,16,128]{2,1,0} %multiply.250.clone.1, f32[3,16,128]{2,1,0} %multiply.249.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.977 = f32[] parameter(1)
  %broadcast.803 = f32[3,16,128]{2,1,0} broadcast(f32[] %param_1.977), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.26 = f32[3,16,128]{2,1,0} divide(f32[3,16,128]{2,1,0} %add.340.clone.1, f32[3,16,128]{2,1,0} %broadcast.803), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.1 = f32[3,16,128]{2,1,0} sqrt(f32[3,16,128]{2,1,0} %divide.26), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_484 = f32[] constant(1e-08)
  %broadcast.802 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_484), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.339 = f32[3,16,128]{2,1,0} add(f32[3,16,128]{2,1,0} %sqrt.1, f32[3,16,128]{2,1,0} %broadcast.802), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.248 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %broadcast.805, f32[3,16,128]{2,1,0} %add.339)
  %divide.25 = f32[3,16,128]{2,1,0} divide(f32[3,16,128]{2,1,0} %add.341.clone.1, f32[3,16,128]{2,1,0} %multiply.248), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.883 = f32[] parameter(3)
  %constant_483 = f32[] constant(-0.02)
  %multiply.398 = f32[] multiply(f32[] %param_3.883, f32[] %constant_483), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.801 = f32[3,16,128]{2,1,0} broadcast(f32[] %multiply.398), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.247 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %divide.25, f32[3,16,128]{2,1,0} %broadcast.801), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.338 = f32[3,16,128]{2,1,0} add(f32[3,16,128]{2,1,0} %param_0.24, f32[3,16,128]{2,1,0} %multiply.247), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.48 = (f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}) tuple(f32[3,16,128]{2,1,0} %add.338, f32[3,16,128]{2,1,0} %add.340.clone.1, f32[3,16,128]{2,1,0} %add.341.clone.1)
}

%region_26.1659 (Arg_0.1660: f32[], Arg_1.1661: f32[]) -> f32[] {
  %Arg_0.1660 = f32[] parameter(0)
  %Arg_1.1661 = f32[] parameter(1)
  ROOT %add.1662 = f32[] add(f32[] %Arg_0.1660, f32[] %Arg_1.1661), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.4 (param_0.506: f32[3,16,4096,221], param_1.678: f32[3,4096,221]) -> f32[3,16,128] {
  %constant_490 = f32[] constant(0)
  %broadcast.965 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_490), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 16, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.111 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_786 = s32[] constant(0)
  %broadcast.1219 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_786), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.678 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.63 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.678), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.42 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.63), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_785 = s32[] constant(127)
  %broadcast.1218 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_785), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.31 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1219, s32[3,4096,221]{2,1,0} %convert.42, s32[3,4096,221]{2,1,0} %broadcast.1218), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.157 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.31, s32[3,4096,221]{2,1,0} %broadcast.1219), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_784 = s32[] constant(128)
  %broadcast.1217 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_784), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.528 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.31, s32[3,4096,221]{2,1,0} %broadcast.1217), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.169 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.157, s32[3,4096,221]{2,1,0} %add.528, s32[3,4096,221]{2,1,0} %clamp.31), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.166 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.169), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.150 = s32[3,4096,221,2]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.111, s32[3,4096,221,1]{2,1,0,3} %bitcast.166), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_772 = f32[] constant(1)
  %broadcast.1207 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_772), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.124 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.678, f32[3,4096,221]{2,1,0} %floor.63), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.123 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1207, f32[3,4096,221]{2,1,0} %subtract.124), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1204 = f32[3,16,4096,221]{2,1,0,3} broadcast(f32[3,4096,221]{2,1,0} %subtract.123), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.506 = f32[3,16,4096,221]{2,1,0,3} parameter(0)
  %multiply.255 = f32[3,16,4096,221]{2,1,0,3} multiply(f32[3,16,4096,221]{2,1,0,3} %broadcast.1204, f32[3,16,4096,221]{2,1,0,3} %param_0.506), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.46 = f32[3,16,128]{2,1,0} scatter(f32[3,16,128]{2,1,0} %broadcast.965, s32[3,4096,221,2]{2,1,0,3} %concatenate.150, f32[3,16,4096,221]{2,1,0,3} %multiply.255), update_window_dims={1}, inserted_window_dims={0,2}, scatter_dims_to_operand_dims={0,2}, index_vector_dim=3, to_apply=%region_26.1659, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_25.1655 (Arg_0.1656: f32[], Arg_1.1657: f32[]) -> f32[] {
  %Arg_0.1656 = f32[] parameter(0)
  %Arg_1.1657 = f32[] parameter(1)
  ROOT %add.1658 = f32[] add(f32[] %Arg_0.1656, f32[] %Arg_1.1657), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.5 (param_0.501: f32[3,16,4096,221], param_1.674: f32[3,4096,221]) -> f32[3,16,128] {
  %constant_491 = f32[] constant(0)
  %broadcast.966 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_491), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 16, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.107 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_760 = s32[] constant(0)
  %broadcast.1197 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_760), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.674 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.55 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.674), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.38 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.55), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_759 = s32[] constant(1)
  %broadcast.1196 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_759), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.523 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.38, s32[3,4096,221]{2,1,0} %broadcast.1196), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_758 = s32[] constant(127)
  %broadcast.1195 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_758), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.27 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1197, s32[3,4096,221]{2,1,0} %add.523, s32[3,4096,221]{2,1,0} %broadcast.1195), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.153 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.27, s32[3,4096,221]{2,1,0} %broadcast.1197), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_757 = s32[] constant(128)
  %broadcast.1194 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_757), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.522 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.27, s32[3,4096,221]{2,1,0} %broadcast.1194), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.165 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.153, s32[3,4096,221]{2,1,0} %add.522, s32[3,4096,221]{2,1,0} %clamp.27), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.162 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.165), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.146 = s32[3,4096,221,2]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.107, s32[3,4096,221,1]{3,2,1,0} %bitcast.162), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.77 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.674, f32[3,4096,221]{2,1,0} %floor.55), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.967 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %subtract.77), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.501 = f32[3,16,4096,221]{2,1,0,3} parameter(0)
  %copy.21 = f32[3,16,4096,221]{3,2,1,0} copy(f32[3,16,4096,221]{2,1,0,3} %param_0.501), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %multiply.256 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.967, f32[3,16,4096,221]{3,2,1,0} %copy.21), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.47 = f32[3,16,128]{2,1,0} scatter(f32[3,16,128]{2,1,0} %broadcast.966, s32[3,4096,221,2]{3,2,1,0} %concatenate.146, f32[3,16,4096,221]{3,2,1,0} %multiply.256), update_window_dims={1}, inserted_window_dims={0,2}, scatter_dims_to_operand_dims={0,2}, index_vector_dim=3, to_apply=%region_25.1655, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%fused_computation.10 (param_0.732: f32[3,16,128,128], param_1.962: f32[3,4096,221], param_2.978: f32[3,4096,221], param_3.849: f32[4096,221], param_4.605: f32[4096,221], param_5.544: pred[4096], param_6.467: f32[4096], param_7.400: f32[4096], param_8.325: f32[4096,221], param_9.201: f32[4096,221], param_10.108: f32[4096,221], param_11.75: f32[4096,110], param_12.47: f32[4096,110]) -> f32[3,16,4096,221] {
  %param_10.108 = f32[4096,221]{1,0} parameter(10)
  %slice.700 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_10.108), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_11.75 = f32[4096,110]{1,0} parameter(11)
  %param_12.47 = f32[4096,110]{1,0} parameter(12)
  %add.993 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %param_11.75, f32[4096,110]{1,0} %param_12.47), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.699 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_10.108), slice={[0:4096:1], [0:219:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.25 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.699), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.992 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %add.993, f32[4096,110]{1,0} %reverse.25), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.253 = f32[4096,111]{1,0} concatenate(f32[4096,1]{1,0} %slice.700, f32[4096,110]{1,0} %add.992), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_1757 = f32[] constant(0)
  %pad.170 = f32[4096,221]{1,0} pad(f32[4096,111]{1,0} %concatenate.253, f32[] %constant_1757), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.169 = f32[4096,221]{1,0} pad(f32[4096,110]{1,0} %add.993, f32[] %constant_1757), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.991 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %pad.170, f32[4096,221]{1,0} %pad.169), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.24 = f32[4096,221]{1,0} reverse(f32[4096,221]{1,0} %add.991), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/rev[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_8.325 = f32[4096,221]{1,0} parameter(8)
  %param_9.201 = f32[4096,221]{1,0} parameter(9)
  %slice.698 = f32[4096,220]{1,0} slice(f32[4096,221]{1,0} %param_9.201), slice={[0:4096], [0:220]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(4096, 220) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=286}
  %constant_1084 = f32[] constant(1)
  %pad.168 = f32[4096,221]{1,0} pad(f32[4096,220]{1,0} %slice.698, f32[] %constant_1084), padding=0_0x1_0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=283}
  %multiply.595 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %param_8.325, f32[4096,221]{1,0} %pad.168), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=280}
  %negate.49 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.595), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %param_3.849 = f32[4096,221]{1,0} parameter(3)
  %negate.48 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_3.849), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_5.544 = pred[4096]{0} parameter(5)
  %param_7.400 = f32[4096]{0} parameter(7)
  %param_6.467 = f32[4096]{0} parameter(6)
  %constant_1758 = f32[] constant(0.001)
  %broadcast.2203 = f32[4096]{0} broadcast(f32[] %constant_1758), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.990 = f32[4096]{0} add(f32[4096]{0} %param_6.467, f32[4096]{0} %broadcast.2203), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.93 = f32[4096]{0} maximum(f32[4096]{0} %param_7.400, f32[4096]{0} %add.990), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.395 = f32[4096]{0} select(pred[4096]{0} %param_5.544, f32[4096]{0} %maximum.93, f32[4096]{0} %broadcast.2203), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %broadcast.2202 = f32[4096]{0} broadcast(f32[] %constant_1757), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.394 = f32[4096]{0} select(pred[4096]{0} %param_5.544, f32[4096]{0} %param_6.467, f32[4096]{0} %broadcast.2202), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.370 = f32[4096]{0} subtract(f32[4096]{0} %select.395, f32[4096]{0} %select.394), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1755 = f32[] constant(0.00452488707)
  %broadcast.2201 = f32[4096]{0} broadcast(f32[] %constant_1755), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.594 = f32[4096]{0} multiply(f32[4096]{0} %subtract.370, f32[4096]{0} %broadcast.2201), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.2200 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.594), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.593 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.48, f32[4096,221]{1,0} %broadcast.2200), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %exponential.28 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %multiply.593), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.592 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.49, f32[4096,221]{1,0} %exponential.28), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %add.989 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %reverse.24, f32[4096,221]{1,0} %multiply.592), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.590 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %add.989, f32[4096,221]{1,0} %broadcast.2200), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %negate.47 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.590), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_4.605 = f32[4096,221]{1,0} parameter(4)
  %constant_1754 = f32[] constant(10)
  %broadcast.2195 = f32[4096,221]{1,0} broadcast(f32[] %constant_1754), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %add.987 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %param_4.605, f32[4096,221]{1,0} %broadcast.2195), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %constant_1753 = f32[] constant(inf)
  %broadcast.2194 = f32[4096,221]{1,0} broadcast(f32[] %constant_1753), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.339 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %add.987, f32[4096,221]{1,0} %broadcast.2194), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2193 = f32[4096,221]{1,0} broadcast(f32[] %constant_1757), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.391 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.339, f32[4096,221]{1,0} %broadcast.2193, f32[4096,221]{1,0} %add.987), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.338 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %param_3.849, f32[4096,221]{1,0} %broadcast.2194), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.390 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.338, f32[4096,221]{1,0} %broadcast.2193, f32[4096,221]{1,0} %param_3.849), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %subtract.368 = f32[4096,221]{1,0} subtract(f32[4096,221]{1,0} %select.391, f32[4096,221]{1,0} %select.390), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %exponential.27 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %subtract.368), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %multiply.589 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.47, f32[4096,221]{1,0} %exponential.27), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(softplus)))/jit(jit_transpose(jvp(softplus)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2192 = f32[3,16,4096,221]{2,1,0,3} broadcast(f32[4096,221]{1,0} %multiply.589), dimensions={2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 16, 4096, 221) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %broadcast.1513 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_1084), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.978 = f32[3,4096,221]{2,1,0} parameter(2)
  %floor.175 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_2.978), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.196 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_2.978, f32[3,4096,221]{2,1,0} %floor.175), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.195 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1513, f32[3,4096,221]{2,1,0} %subtract.196), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.962 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.174 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.962), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.194 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.962, f32[3,4096,221]{2,1,0} %floor.174), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.193 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1513, f32[3,4096,221]{2,1,0} %subtract.194), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.457 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.195, f32[3,4096,221]{2,1,0} %subtract.193), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1512 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.457), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.732 = f32[3,16,128,128]{3,2,1,0} parameter(0)
  %iota.135 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1072 = s32[] constant(0)
  %broadcast.1511 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1072), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.106 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.175), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1070 = s32[] constant(127)
  %broadcast.1510 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1070), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.95 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1511, s32[3,4096,221]{2,1,0} %convert.106, s32[3,4096,221]{2,1,0} %broadcast.1510), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.221 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.95, s32[3,4096,221]{2,1,0} %broadcast.1511), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1069 = s32[] constant(128)
  %broadcast.1509 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1069), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.651 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.95, s32[3,4096,221]{2,1,0} %broadcast.1509), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.233 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.221, s32[3,4096,221]{2,1,0} %add.651, s32[3,4096,221]{2,1,0} %clamp.95), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.222 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.233), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.105 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.174), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.94 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1511, s32[3,4096,221]{2,1,0} %convert.105, s32[3,4096,221]{2,1,0} %broadcast.1510), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.220 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.94, s32[3,4096,221]{2,1,0} %broadcast.1511), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.650 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.94, s32[3,4096,221]{2,1,0} %broadcast.1509), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.232 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.220, s32[3,4096,221]{2,1,0} %add.650, s32[3,4096,221]{2,1,0} %clamp.94), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.221 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.232), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.174 = s32[3,4096,221,3]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.135, s32[3,4096,221,1]{2,1,0,3} %bitcast.222, s32[3,4096,221,1]{2,1,0,3} %bitcast.221), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.87 = f32[3,16,4096,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_0.732, s32[3,4096,221,3]{2,1,0,3} %concatenate.174), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.456 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.1512, f32[3,16,4096,221]{3,2,1,0} %gather.87), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.455 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.195, f32[3,4096,221]{2,1,0} %subtract.194), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1504 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.455), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.134 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.220 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.233), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1071 = s32[] constant(1)
  %broadcast.1499 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1071), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.648 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.105, s32[3,4096,221]{2,1,0} %broadcast.1499), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.92 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1511, s32[3,4096,221]{2,1,0} %add.648, s32[3,4096,221]{2,1,0} %broadcast.1510), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.218 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.92, s32[3,4096,221]{2,1,0} %broadcast.1511), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.647 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.92, s32[3,4096,221]{2,1,0} %broadcast.1509), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.230 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.218, s32[3,4096,221]{2,1,0} %add.647, s32[3,4096,221]{2,1,0} %clamp.92), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.219 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.230), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.173 = s32[3,4096,221,3]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.134, s32[3,4096,221,1]{3,2,1,0} %bitcast.220, s32[3,4096,221,1]{3,2,1,0} %bitcast.219), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.86 = f32[3,16,4096,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_0.732, s32[3,4096,221,3]{3,2,1,0} %concatenate.173), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.454 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.1504, f32[3,16,4096,221]{3,2,1,0} %gather.86), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.646 = f32[3,16,4096,221]{3,2,1,0} add(f32[3,16,4096,221]{3,2,1,0} %multiply.456, f32[3,16,4096,221]{3,2,1,0} %multiply.454), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.453 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.196, f32[3,4096,221]{2,1,0} %subtract.193), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1495 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.453), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.133 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.645 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.106, s32[3,4096,221]{2,1,0} %broadcast.1499), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.91 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1511, s32[3,4096,221]{2,1,0} %add.645, s32[3,4096,221]{2,1,0} %broadcast.1510), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.217 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.91, s32[3,4096,221]{2,1,0} %broadcast.1511), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.644 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.91, s32[3,4096,221]{2,1,0} %broadcast.1509), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.229 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.217, s32[3,4096,221]{2,1,0} %add.644, s32[3,4096,221]{2,1,0} %clamp.91), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.218 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.229), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.172 = s32[3,4096,221,3]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.133, s32[3,4096,221,1]{2,1,0,3} %bitcast.218, s32[3,4096,221,1]{2,1,0,3} %bitcast.221), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.85 = f32[3,16,4096,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_0.732, s32[3,4096,221,3]{2,1,0,3} %concatenate.172), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.452 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.1495, f32[3,16,4096,221]{3,2,1,0} %gather.85), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.642 = f32[3,16,4096,221]{3,2,1,0} add(f32[3,16,4096,221]{3,2,1,0} %add.646, f32[3,16,4096,221]{3,2,1,0} %multiply.452), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.451 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.196, f32[3,4096,221]{2,1,0} %subtract.194), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1487 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.451), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.132 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.216 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.229), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.171 = s32[3,4096,221,3]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.132, s32[3,4096,221,1]{3,2,1,0} %bitcast.216, s32[3,4096,221,1]{3,2,1,0} %bitcast.219), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.84 = f32[3,16,4096,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_0.732, s32[3,4096,221,3]{3,2,1,0} %concatenate.171), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.450 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.1487, f32[3,16,4096,221]{3,2,1,0} %gather.84), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.637 = f32[3,16,4096,221]{3,2,1,0} add(f32[3,16,4096,221]{3,2,1,0} %add.642, f32[3,16,4096,221]{3,2,1,0} %multiply.450), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.22 = f32[3,16,4096,221]{2,1,0,3} copy(f32[3,16,4096,221]{3,2,1,0} %add.637), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %multiply.257 = f32[3,16,4096,221]{2,1,0,3} multiply(f32[3,16,4096,221]{2,1,0,3} %broadcast.2192, f32[3,16,4096,221]{2,1,0,3} %copy.22), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
}

%fused_computation.16 (param_0.717: f32[4096,3], param_1.936: f32[4096,13]) -> f32[4096,6] {
  %param_1.936 = f32[4096,13]{1,0} parameter(1)
  %slice.674 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_1.936), slice={[0:4096:1], [0:12:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 12) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.673 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_1.936), slice={[0:4096:1], [1:13:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 13) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.952 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %slice.674, f32[4096,6]{1,0} %slice.673), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.288 = f32[4096,1]{1,0} slice(f32[4096,6]{1,0} %add.952), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_0.717 = f32[4096,3]{1,0} parameter(0)
  %slice.287 = f32[4096,2]{1,0} slice(f32[4096,3]{1,0} %param_0.717), slice={[0:4096], [0:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 2) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.286 = f32[4096,2]{1,0} slice(f32[4096,6]{1,0} %add.952), slice={[0:4096:1], [2:6:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 6) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.354 = f32[4096,2]{1,0} add(f32[4096,2]{1,0} %slice.287, f32[4096,2]{1,0} %slice.286), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.106 = f32[4096,3]{1,0} concatenate(f32[4096,1]{1,0} %slice.288, f32[4096,2]{1,0} %add.354), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_498 = f32[] constant(0)
  ROOT %pad.104 = f32[4096,6]{1,0} pad(f32[4096,3]{1,0} %concatenate.106, f32[] %constant_498), padding=0_0_0x0_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.17 (param_0.720: f32[4096,13]) -> f32[4096,3] {
  %param_0.720 = f32[4096,13]{1,0} parameter(0)
  %slice.678 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_0.720), slice={[0:4096:1], [0:12:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 12) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.677 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_0.720), slice={[0:4096:1], [1:13:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 13) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.955 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %slice.678, f32[4096,6]{1,0} %slice.677), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.294 = f32[4096,3]{1,0} slice(f32[4096,6]{1,0} %add.955), slice={[0:4096:1], [0:5:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 5) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.293 = f32[4096,3]{1,0} slice(f32[4096,6]{1,0} %add.955), slice={[0:4096:1], [1:6:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 6) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.358 = f32[4096,3]{1,0} add(f32[4096,3]{1,0} %slice.294, f32[4096,3]{1,0} %slice.293), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.292 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.358), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.291 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.358), slice={[0:4096:1], [0:2:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 2) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.290 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.358), slice={[0:4096:1], [1:3:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 3) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.357 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %slice.291, f32[4096,1]{1,0} %slice.290), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.289 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.358), slice={[0:4096:1], [2:3:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 3) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.356 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %add.357, f32[4096,1]{1,0} %slice.289), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.107 = f32[4096,2]{1,0} concatenate(f32[4096,1]{1,0} %slice.292, f32[4096,1]{1,0} %add.356), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_499 = f32[] constant(0)
  %pad.106 = f32[4096,3]{1,0} pad(f32[4096,2]{1,0} %concatenate.107, f32[] %constant_499), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.105 = f32[4096,3]{1,0} pad(f32[4096,1]{1,0} %add.357, f32[] %constant_499), padding=0_0x1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %add.355 = f32[4096,3]{1,0} add(f32[4096,3]{1,0} %pad.106, f32[4096,3]{1,0} %pad.105), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.19 (param_0.715: f32[4096,55]) -> f32[4096,13] {
  %param_0.715 = f32[4096,55]{1,0} parameter(0)
  %slice.670 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_0.715), slice={[0:4096:1], [0:54:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 54) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.669 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_0.715), slice={[0:4096:1], [1:55:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 55) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.949 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %slice.670, f32[4096,27]{1,0} %slice.669), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.298 = f32[4096,13]{1,0} slice(f32[4096,27]{1,0} %add.949), slice={[0:4096:1], [0:26:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 26) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.297 = f32[4096,13]{1,0} slice(f32[4096,27]{1,0} %add.949), slice={[0:4096:1], [1:27:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 27) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %add.360 = f32[4096,13]{1,0} add(f32[4096,13]{1,0} %slice.298, f32[4096,13]{1,0} %slice.297), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.21 (param_0.711: f32[4096,221]) -> f32[4096,55] {
  %param_0.711 = f32[4096,221]{1,0} parameter(0)
  %slice.662 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_0.711), slice={[0:4096:1], [2:221:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 220) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.21 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.662), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 220) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.661 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_0.711), slice={[0:4096:1], [1:220:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.20 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.661), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.944 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %reverse.21, f32[4096,110]{1,0} %reverse.20), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.302 = f32[4096,55]{1,0} slice(f32[4096,110]{1,0} %add.944), slice={[0:4096:1], [0:109:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 109) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.301 = f32[4096,55]{1,0} slice(f32[4096,110]{1,0} %add.944), slice={[0:4096:1], [1:110:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 110) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %add.362 = f32[4096,55]{1,0} add(f32[4096,55]{1,0} %slice.302, f32[4096,55]{1,0} %slice.301), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%region_11.1478 (Arg_0.1479: f32[], Arg_1.1480: f32[]) -> f32[] {
  %Arg_0.1479 = f32[] parameter(0)
  %Arg_1.1480 = f32[] parameter(1)
  ROOT %add.1481 = f32[] add(f32[] %Arg_0.1479, f32[] %Arg_1.1480), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=196}
}

%fused_computation.23 (param_0.61: f32[4096,221], param_1.778: f32[4096,221], param_2.742: f32[4096], param_3.817: pred[4096], param_4.593: f32[4096], param_5.534: f32[4096], param_6.455: f32[4096,221], param_7.384: f32[4096], param_8.315: f32[4096,3]) -> f32[4096,221] {
  %param_8.315 = f32[4096,3]{1,0} parameter(8)
  %constant_1723 = f32[] constant(0.000162760422)
  %broadcast.2179 = f32[4096,3]{1,0} broadcast(f32[] %constant_1723), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %multiply.581 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %param_8.315, f32[4096,3]{1,0} %broadcast.2179), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %constant_500 = f32[] constant(0)
  %reduce.37 = f32[4096]{0} reduce(f32[4096,3]{1,0} %multiply.581, f32[] %constant_500), dimensions={1}, to_apply=%region_11.1478, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=196}
  %bitcast.78 = f32[4096,1]{1,0} bitcast(f32[4096]{0} %reduce.37)
  %pad.109 = f32[4096,221]{1,0} pad(f32[4096,1]{1,0} %bitcast.78, f32[] %constant_500), padding=0_0x220_0, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0, 1), inserted_window_dims=(), scatter_dims_to_operand_dims=(1,)) indices_are_sorted=True unique_indices=True mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=196}
  %param_2.742 = f32[4096]{0} parameter(2)
  %param_7.384 = f32[4096]{0} parameter(7)
  %constant_1709 = f32[] constant(1e-08)
  %broadcast.2159 = f32[4096]{0} broadcast(f32[] %constant_1709), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.940 = f32[4096]{0} add(f32[4096]{0} %param_7.384, f32[4096]{0} %broadcast.2159), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  %divide.27 = f32[4096]{0} divide(f32[4096]{0} %param_2.742, f32[4096]{0} %add.940), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %negate.18 = f32[4096]{0} negate(f32[4096]{0} %divide.27), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %bitcast.77 = f32[4096,1]{1,0} bitcast(f32[4096]{0} %negate.18)
  %pad.108 = f32[4096,221]{1,0} pad(f32[4096,1]{1,0} %bitcast.77, f32[] %constant_500), padding=0_0x220_0, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0,), inserted_window_dims=(1,), scatter_dims_to_operand_dims=(1,)) indices_are_sorted=True unique_indices=True mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=482}
  %add.365 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %pad.109, f32[4096,221]{1,0} %pad.108), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=482}
  %constant_1214 = f32[] constant(1)
  %broadcast.1646 = f32[4096,221]{1,0} broadcast(f32[] %constant_1214), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %param_6.455 = f32[4096,221]{1,0} parameter(6)
  %negate.35 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_6.455), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_3.817 = pred[4096]{0} parameter(3)
  %param_5.534 = f32[4096]{0} parameter(5)
  %param_4.593 = f32[4096]{0} parameter(4)
  %constant_1213 = f32[] constant(0.001)
  %broadcast.1644 = f32[4096]{0} broadcast(f32[] %constant_1213), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.698 = f32[4096]{0} add(f32[4096]{0} %param_4.593, f32[4096]{0} %broadcast.1644), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.85 = f32[4096]{0} maximum(f32[4096]{0} %param_5.534, f32[4096]{0} %add.698), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.275 = f32[4096]{0} select(pred[4096]{0} %param_3.817, f32[4096]{0} %maximum.85, f32[4096]{0} %broadcast.1644), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %broadcast.1642 = f32[4096]{0} broadcast(f32[] %constant_500), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.274 = f32[4096]{0} select(pred[4096]{0} %param_3.817, f32[4096]{0} %param_4.593, f32[4096]{0} %broadcast.1642), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.236 = f32[4096]{0} subtract(f32[4096]{0} %select.275, f32[4096]{0} %select.274), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1211 = f32[] constant(0.00452488707)
  %broadcast.1641 = f32[4096]{0} broadcast(f32[] %constant_1211), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.499 = f32[4096]{0} multiply(f32[4096]{0} %subtract.236, f32[4096]{0} %broadcast.1641), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.1640 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.499), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.498 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.35, f32[4096,221]{1,0} %broadcast.1640), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %exponential.20 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %multiply.498), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %subtract.235 = f32[4096,221]{1,0} subtract(f32[4096,221]{1,0} %broadcast.1646, f32[4096,221]{1,0} %exponential.20), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %param_1.778 = f32[4096,221]{1,0} parameter(1)
  %multiply.263 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %subtract.235, f32[4096,221]{1,0} %param_1.778), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=280}
  %slice.305 = f32[4096,220]{1,0} slice(f32[4096,221]{1,0} %multiply.263), slice={[0:4096], [1:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=283}
  %pad.107 = f32[4096,221]{1,0} pad(f32[4096,220]{1,0} %slice.305, f32[] %constant_500), padding=0_0x0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0, 1), inserted_window_dims=(), scatter_dims_to_operand_dims=(1,)) indices_are_sorted=True unique_indices=True mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=286}
  %add.364 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %add.365, f32[4096,221]{1,0} %pad.107), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=286}
  %param_0.61 = f32[4096,221]{1,0} parameter(0)
  ROOT %multiply.262 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %add.364, f32[4096,221]{1,0} %param_0.61), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%region_16.1498 (Arg_0.1499: f32[], Arg_1.1500: f32[]) -> f32[] {
  %Arg_0.1499 = f32[] parameter(0)
  %Arg_1.1500 = f32[] parameter(1)
  ROOT %add.1501 = f32[] add(f32[] %Arg_0.1499, f32[] %Arg_1.1500), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.6 (param_0.62: f32[4096,221], param_1.88: s32[4096,33,2], param_2.922: f32[4096], param_3.812: f32[4096,221], param_4.589: f32[4096]) -> f32[4096,221] {
  %param_0.62 = f32[4096,221]{1,0} parameter(0)
  %param_1.88 = s32[4096,33,2]{1,0,2} parameter(1)
  %param_2.922 = f32[4096]{0} parameter(2)
  %constant_501 = f32[] constant(1)
  %broadcast.815 = f32[4096]{0} broadcast(f32[] %constant_501), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_4.589 = f32[4096]{0} parameter(4)
  %constant_1699 = f32[] constant(1e-08)
  %broadcast.2155 = f32[4096]{0} broadcast(f32[] %constant_1699), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.938 = f32[4096]{0} add(f32[4096]{0} %param_4.589, f32[4096]{0} %broadcast.2155), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  %multiply.266 = f32[4096]{0} multiply(f32[4096]{0} %add.938, f32[4096]{0} %add.938), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %divide.28 = f32[4096]{0} divide(f32[4096]{0} %broadcast.815, f32[4096]{0} %multiply.266), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %multiply.265 = f32[4096]{0} multiply(f32[4096]{0} %param_2.922, f32[4096]{0} %divide.28), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %broadcast.2147 = f32[4096,1]{1,0} broadcast(f32[] %constant_501), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_3.812 = f32[4096,221]{1,0} parameter(3)
  %slice.654 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_3.812), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(1,), start_index_map=(1,)) slice_sizes=(4096, 1) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=482}
  %subtract.362 = f32[4096,1]{1,0} subtract(f32[4096,1]{1,0} %broadcast.2147, f32[4096,1]{1,0} %slice.654), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %broadcast.2145 = f32[4096,1]{1,0} broadcast(f32[] %constant_1699), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.933 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %subtract.362, f32[4096,1]{1,0} %broadcast.2145), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %bitcast.544 = f32[4096]{0} bitcast(f32[4096,1]{1,0} %add.933), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %multiply.264 = f32[4096]{0} multiply(f32[4096]{0} %multiply.265, f32[4096]{0} %bitcast.544), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %negate.19 = f32[4096]{0} negate(f32[4096]{0} %multiply.264), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %broadcast.814 = f32[4096,33]{1,0} broadcast(f32[4096]{0} %negate.19), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/broadcast_in_dim[shape=(4096, 33) broadcast_dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  ROOT %scatter.48 = f32[4096,221]{1,0} scatter(f32[4096,221]{1,0} %param_0.62, s32[4096,33,2]{1,0,2} %param_1.88, f32[4096,33]{1,0} %broadcast.814), update_window_dims={}, inserted_window_dims={0,1}, scatter_dims_to_operand_dims={0,1}, index_vector_dim=2, to_apply=%region_16.1498, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(), inserted_window_dims=(0, 1), scatter_dims_to_operand_dims=(0, 1)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=472}
}

%region_18.1506 (Arg_0.1507: f32[], Arg_1.1508: f32[]) -> f32[] {
  %Arg_0.1507 = f32[] parameter(0)
  %Arg_1.1508 = f32[] parameter(1)
  ROOT %add.1509 = f32[] add(f32[] %Arg_0.1507, f32[] %Arg_1.1508), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
}

%fused_computation.24 (param_0.378: f32[4096,221,3], param_1.920: f32[4096,3], param_2.1057: f32[4096], param_3.938: f32[4096,221]) -> f32[4096,221] {
  %param_0.378 = f32[4096,221,3]{1,0,2} parameter(0)
  %param_1.920 = f32[4096,3]{1,0} parameter(1)
  %constant_1894 = f32[] constant(1)
  %broadcast.2352 = f32[4096,1]{1,0} broadcast(f32[] %constant_1894), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_3.938 = f32[4096,221]{1,0} parameter(3)
  %slice.736 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_3.938), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(1,), start_index_map=(1,)) slice_sizes=(4096, 1) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=482}
  %subtract.408 = f32[4096,1]{1,0} subtract(f32[4096,1]{1,0} %broadcast.2352, f32[4096,1]{1,0} %slice.736), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %constant_1893 = f32[] constant(1e-08)
  %broadcast.2351 = f32[4096,1]{1,0} broadcast(f32[] %constant_1893), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.1077 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %subtract.408, f32[4096,1]{1,0} %broadcast.2351), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %bitcast.554 = f32[4096]{0} bitcast(f32[4096,1]{1,0} %add.1077), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_2.1057 = f32[4096]{0} parameter(2)
  %broadcast.2350 = f32[4096]{0} broadcast(f32[] %constant_1893), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.1076 = f32[4096]{0} add(f32[4096]{0} %param_2.1057, f32[4096]{0} %broadcast.2350), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  %divide.64 = f32[4096]{0} divide(f32[4096]{0} %bitcast.554, f32[4096]{0} %add.1076), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %constant_1712 = f32[] constant(0.000162760422)
  %broadcast.2169 = f32[4096]{0} broadcast(f32[] %constant_1712), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.573 = f32[4096]{0} multiply(f32[4096]{0} %divide.64, f32[4096]{0} %broadcast.2169), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %broadcast.2168 = f32[4096,3]{1,0} broadcast(f32[4096]{0} %multiply.573), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.572 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %param_1.920, f32[4096,3]{1,0} %broadcast.2168), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %broadcast.2167 = f32[4096,221,3]{1,0,2} broadcast(f32[4096,3]{1,0} %multiply.572), dimensions={0,2}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/broadcast_in_dim[shape=(4096, 221, 3) broadcast_dimensions=(0, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.267 = f32[4096,221,3]{1,0,2} multiply(f32[4096,221,3]{1,0,2} %param_0.378, f32[4096,221,3]{1,0,2} %broadcast.2167), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %bitcast.79 = f32[3,4096,221]{2,1,0} bitcast(f32[4096,221,3]{1,0,2} %multiply.267), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %constant_504 = f32[] constant(0)
  ROOT %reduce.38 = f32[4096,221]{1,0} reduce(f32[3,4096,221]{2,1,0} %bitcast.79, f32[] %constant_504), dimensions={0}, to_apply=%region_18.1506, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
}

%region_14.1490 (Arg_0.1491: f32[], Arg_1.1492: f32[]) -> f32[] {
  %Arg_0.1491 = f32[] parameter(0)
  %Arg_1.1492 = f32[] parameter(1)
  ROOT %add.1493 = f32[] add(f32[] %Arg_0.1491, f32[] %Arg_1.1492), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
}

%fused_computation.25 (param_0.705: f32[3,4096], param_1.924: f32[4096,3]) -> f32[4096] {
  %param_0.705 = f32[3,4096]{0,1} parameter(0)
  %bitcast.139 = f32[4096,3]{1,0} bitcast(f32[3,4096]{0,1} %param_0.705), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %param_1.924 = f32[4096,3]{1,0} parameter(1)
  %constant_1720 = f32[] constant(0.000162760422)
  %broadcast.2177 = f32[4096,3]{1,0} broadcast(f32[] %constant_1720), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %multiply.579 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %param_1.924, f32[4096,3]{1,0} %broadcast.2177), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %multiply.268 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %bitcast.139, f32[4096,3]{1,0} %multiply.579), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %constant_502 = f32[] constant(0)
  ROOT %reduce.39 = f32[4096]{0} reduce(f32[4096,3]{1,0} %multiply.268, f32[] %constant_502), dimensions={1}, to_apply=%region_14.1490, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
}

%fused_computation.27 (param_0.68: f32[3,48,128,128], param_1.978: f32[], param_2.1000: f32[], param_3.885: f32[], param_4.659: f32[3,48,128,128], param_5.619: f32[3,48,128,128], param_6.553: f32[3,48,128,128], param_7.503: f32[3,48,128,128]) -> (f32[3,48,128,128], f32[3,48,128,128], f32[3,48,128,128]) {
  %param_0.68 = f32[3,48,128,128]{3,2,1,0} parameter(0)
  %param_5.619 = f32[3,48,128,128]{3,2,1,0} parameter(5)
  %param_6.553 = f32[3,48,128,128]{3,2,1,0} parameter(6)
  %add.495.clone.1 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %param_5.619, f32[3,48,128,128]{3,2,1,0} %param_6.553), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_510_clone_1 = f32[] constant(0.0100000007)
  %broadcast.828.clone.1 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_510_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.277.clone.1 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %add.495.clone.1, f32[3,48,128,128]{3,2,1,0} %broadcast.828.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_7.503 = f32[3,48,128,128]{3,2,1,0} parameter(7)
  %constant_511_clone_1 = f32[] constant(0.9)
  %broadcast.827.clone.1 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_511_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.276.clone.1 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %param_7.503, f32[3,48,128,128]{3,2,1,0} %broadcast.827.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.370.clone.1 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %multiply.277.clone.1, f32[3,48,128,128]{3,2,1,0} %multiply.276.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.1000 = f32[] parameter(2)
  %broadcast.821 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %param_2.1000), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_509_clone_1 = f32[] constant(0.1)
  %broadcast.826.clone.1 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_509_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.275.clone.1 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %add.495.clone.1, f32[3,48,128,128]{3,2,1,0} %broadcast.826.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.274.clone.1 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %multiply.275.clone.1, f32[3,48,128,128]{3,2,1,0} %multiply.275.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_507_clone_1 = f32[] constant(0.01)
  %broadcast.825.clone.1 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_507_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.273.clone.1 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %multiply.274.clone.1, f32[3,48,128,128]{3,2,1,0} %broadcast.825.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.659 = f32[3,48,128,128]{3,2,1,0} parameter(4)
  %constant_508_clone_1 = f32[] constant(0.99)
  %broadcast.823.clone.1 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_508_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.272.clone.1 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %param_4.659, f32[3,48,128,128]{3,2,1,0} %broadcast.823.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.369.clone.1 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %multiply.273.clone.1, f32[3,48,128,128]{3,2,1,0} %multiply.272.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.978 = f32[] parameter(1)
  %broadcast.820 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %param_1.978), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.30 = f32[3,48,128,128]{3,2,1,0} divide(f32[3,48,128,128]{3,2,1,0} %add.369.clone.1, f32[3,48,128,128]{3,2,1,0} %broadcast.820), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.2 = f32[3,48,128,128]{3,2,1,0} sqrt(f32[3,48,128,128]{3,2,1,0} %divide.30), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_506 = f32[] constant(1e-08)
  %broadcast.819 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_506), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.368 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %sqrt.2, f32[3,48,128,128]{3,2,1,0} %broadcast.819), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.271 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %broadcast.821, f32[3,48,128,128]{3,2,1,0} %add.368)
  %divide.29 = f32[3,48,128,128]{3,2,1,0} divide(f32[3,48,128,128]{3,2,1,0} %add.370.clone.1, f32[3,48,128,128]{3,2,1,0} %multiply.271), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.885 = f32[] parameter(3)
  %constant_505 = f32[] constant(-0.02)
  %multiply.399 = f32[] multiply(f32[] %param_3.885, f32[] %constant_505), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.818 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %multiply.399), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.270 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %divide.29, f32[3,48,128,128]{3,2,1,0} %broadcast.818), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.366 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %param_0.68, f32[3,48,128,128]{3,2,1,0} %multiply.270), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.52 = (f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}) tuple(f32[3,48,128,128]{3,2,1,0} %add.366, f32[3,48,128,128]{3,2,1,0} %add.369.clone.1, f32[3,48,128,128]{3,2,1,0} %add.370.clone.1)
}

%region_35.1742 (Arg_0.1743: f32[], Arg_1.1744: f32[]) -> f32[] {
  %Arg_0.1743 = f32[] parameter(0)
  %Arg_1.1744 = f32[] parameter(1)
  ROOT %add.1745 = f32[] add(f32[] %Arg_0.1743, f32[] %Arg_1.1744), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.7 (param_0.680: f32[3,48,4096,33], param_1.900: f32[3,3,4096,33]) -> f32[3,48,128,128] {
  %constant_512 = f32[] constant(0)
  %broadcast.969 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_512), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 48, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.175 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1543 = s32[] constant(0)
  %broadcast.1969 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1543), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.900 = f32[3,3,4096,33]{3,2,0,1} parameter(1)
  %slice.570 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.900), slice={[0:3], [1:2], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 1, 4096, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.440 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.570), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.316 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.440), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.186 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.316), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1542 = s32[] constant(127)
  %broadcast.1968 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1542), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.175 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1969, s32[3,4096,33]{2,1,0} %convert.186, s32[3,4096,33]{2,1,0} %broadcast.1968), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.303 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.175, s32[3,4096,33]{2,1,0} %broadcast.1969), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1541 = s32[] constant(128)
  %broadcast.1967 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1541), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.851 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.175, s32[3,4096,33]{2,1,0} %broadcast.1967), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.349 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.303, s32[3,4096,33]{2,1,0} %add.851, s32[3,4096,33]{2,1,0} %clamp.175), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.439 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.349), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.567 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.900), slice={[0:3], [2:3], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 1, 0, 0) limit_indices=(3, 2, 4096, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.438 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.567), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.315 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.438), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.185 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.315), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.174 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1969, s32[3,4096,33]{2,1,0} %convert.185, s32[3,4096,33]{2,1,0} %broadcast.1968), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.302 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.174, s32[3,4096,33]{2,1,0} %broadcast.1969), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.850 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.174, s32[3,4096,33]{2,1,0} %broadcast.1967), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.348 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.302, s32[3,4096,33]{2,1,0} %add.850, s32[3,4096,33]{2,1,0} %clamp.174), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.437 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.348), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.227 = s32[3,4096,33,3]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.175, s32[3,4096,33,1]{2,1,0,3} %bitcast.439, s32[3,4096,33,1]{2,1,0,3} %bitcast.437), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.23 = s32[3,4096,33,3]{3,2,1,0} copy(s32[3,4096,33,3]{2,1,0,3} %concatenate.227), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1519 = f32[] constant(1)
  %broadcast.1942 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1519), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.310 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.440, f32[3,4096,33]{2,1,0} %floor.316), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.309 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1942, f32[3,4096,33]{2,1,0} %subtract.310), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.308 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.438, f32[3,4096,33]{2,1,0} %floor.315), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.307 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1942, f32[3,4096,33]{2,1,0} %subtract.308), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.532 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.309, f32[3,4096,33]{2,1,0} %subtract.307), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1940 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.532), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.680 = f32[3,48,4096,33]{3,2,1,0} parameter(0)
  %multiply.278 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.1940, f32[3,48,4096,33]{3,2,1,0} %param_0.680), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.49 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.969, s32[3,4096,33,3]{3,2,1,0} %copy.23, f32[3,48,4096,33]{3,2,1,0} %multiply.278), update_window_dims={1}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_35.1742, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_34.1738 (Arg_0.1739: f32[], Arg_1.1740: f32[]) -> f32[] {
  %Arg_0.1739 = f32[] parameter(0)
  %Arg_1.1740 = f32[] parameter(1)
  ROOT %add.1741 = f32[] add(f32[] %Arg_0.1739, f32[] %Arg_1.1740), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.8 (param_0.677: f32[3,48,4096,33], param_1.896: f32[3,48,128,128], param_2.892: f32[3,48,128,128], param_3.758: f32[3,3,4096,33]) -> f32[3,48,128,128] {
  %param_1.896 = f32[3,48,128,128]{3,2,1,0} parameter(1)
  %param_2.892 = f32[3,48,128,128]{3,2,1,0} parameter(2)
  %add.371 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %param_1.896, f32[3,48,128,128]{3,2,1,0} %param_2.892), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.171 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1510 = s32[] constant(0)
  %broadcast.1931 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1510), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_3.758 = f32[3,3,4096,33]{3,2,0,1} parameter(3)
  %slice.540 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_3.758), slice={[0:3], [1:2], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 1, 4096, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.416 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.540), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.300 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.416), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.178 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.300), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1508 = s32[] constant(127)
  %broadcast.1929 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1508), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.167 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1931, s32[3,4096,33]{2,1,0} %convert.178, s32[3,4096,33]{2,1,0} %broadcast.1929), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.294 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.167, s32[3,4096,33]{2,1,0} %broadcast.1931), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1507 = s32[] constant(128)
  %broadcast.1928 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1507), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.841 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.167, s32[3,4096,33]{2,1,0} %broadcast.1928), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.341 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.294, s32[3,4096,33]{2,1,0} %add.841, s32[3,4096,33]{2,1,0} %clamp.167), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.415 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.341), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.539 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_3.758), slice={[0:3], [2:3], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 1, 0, 0) limit_indices=(3, 2, 4096, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.414 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.539), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.299 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.414), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.177 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.299), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1509 = s32[] constant(1)
  %broadcast.1926 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1509), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.840 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.177, s32[3,4096,33]{2,1,0} %broadcast.1926), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.166 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1931, s32[3,4096,33]{2,1,0} %add.840, s32[3,4096,33]{2,1,0} %broadcast.1929), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.293 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.166, s32[3,4096,33]{2,1,0} %broadcast.1931), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.839 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.166, s32[3,4096,33]{2,1,0} %broadcast.1928), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.340 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.293, s32[3,4096,33]{2,1,0} %add.839, s32[3,4096,33]{2,1,0} %clamp.166), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.413 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.340), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.223 = s32[3,4096,33,3]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.171, s32[3,4096,33,1]{3,2,1,0} %bitcast.415, s32[3,4096,33,1]{3,2,1,0} %bitcast.413), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1469 = f32[] constant(1)
  %broadcast.1883 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1469), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.294 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.416, f32[3,4096,33]{2,1,0} %floor.300), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.293 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1883, f32[3,4096,33]{2,1,0} %subtract.294), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.292 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.414, f32[3,4096,33]{2,1,0} %floor.299), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.528 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.293, f32[3,4096,33]{2,1,0} %subtract.292), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1882 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.528), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.677 = f32[3,48,4096,33]{3,2,1,0} parameter(0)
  %multiply.279 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.1882, f32[3,48,4096,33]{3,2,1,0} %param_0.677), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.50 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %add.371, s32[3,4096,33,3]{3,2,1,0} %concatenate.223, f32[3,48,4096,33]{3,2,1,0} %multiply.279), update_window_dims={1}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_34.1738, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_33.1734 (Arg_0.1735: f32[], Arg_1.1736: f32[]) -> f32[] {
  %Arg_0.1735 = f32[] parameter(0)
  %Arg_1.1736 = f32[] parameter(1)
  ROOT %add.1737 = f32[] add(f32[] %Arg_0.1735, f32[] %Arg_1.1736), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.9 (param_0.662: f32[3,48,4096,33], param_1.884: f32[3,3,4096,33]) -> f32[3,48,128,128] {
  %constant_513 = f32[] constant(0)
  %broadcast.970 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_513), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 48, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.167 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1459 = s32[] constant(0)
  %broadcast.1875 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1459), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.884 = f32[3,3,4096,33]{3,2,0,1} parameter(1)
  %slice.520 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.884), slice={[0:3], [1:2], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 1, 4096, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.388 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.520), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.280 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.388), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.166 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.280), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1462 = s32[] constant(1)
  %broadcast.1874 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1462), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.823 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.166, s32[3,4096,33]{2,1,0} %broadcast.1874), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1458 = s32[] constant(127)
  %broadcast.1873 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1458), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.155 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1875, s32[3,4096,33]{2,1,0} %add.823, s32[3,4096,33]{2,1,0} %broadcast.1873), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.281 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.155, s32[3,4096,33]{2,1,0} %broadcast.1875), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1457 = s32[] constant(128)
  %broadcast.1872 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1457), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.822 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.155, s32[3,4096,33]{2,1,0} %broadcast.1872), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.329 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.281, s32[3,4096,33]{2,1,0} %add.822, s32[3,4096,33]{2,1,0} %clamp.155), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.387 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.329), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.519 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.884), slice={[0:3], [2:3], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 1, 0, 0) limit_indices=(3, 2, 4096, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.386 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.519), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.278 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.386), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.165 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.278), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.154 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1875, s32[3,4096,33]{2,1,0} %convert.165, s32[3,4096,33]{2,1,0} %broadcast.1873), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.280 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.154, s32[3,4096,33]{2,1,0} %broadcast.1875), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.821 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.154, s32[3,4096,33]{2,1,0} %broadcast.1872), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.328 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.280, s32[3,4096,33]{2,1,0} %add.821, s32[3,4096,33]{2,1,0} %clamp.154), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.385 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.328), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.219 = s32[3,4096,33,3]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.167, s32[3,4096,33,1]{2,1,0,3} %bitcast.387, s32[3,4096,33,1]{2,1,0,3} %bitcast.385), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.24 = s32[3,4096,33,3]{3,2,1,0} copy(s32[3,4096,33,3]{2,1,0,3} %concatenate.219), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.282 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.388, f32[3,4096,33]{2,1,0} %floor.280), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1419 = f32[] constant(1)
  %broadcast.1835 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1419), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.279 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.386, f32[3,4096,33]{2,1,0} %floor.278), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.278 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1835, f32[3,4096,33]{2,1,0} %subtract.279), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.524 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.282, f32[3,4096,33]{2,1,0} %subtract.278), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1834 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.524), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.662 = f32[3,48,4096,33]{3,2,1,0} parameter(0)
  %multiply.280 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.1834, f32[3,48,4096,33]{3,2,1,0} %param_0.662), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.51 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.970, s32[3,4096,33,3]{3,2,1,0} %copy.24, f32[3,48,4096,33]{3,2,1,0} %multiply.280), update_window_dims={1}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_33.1734, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_32.1730 (Arg_0.1731: f32[], Arg_1.1732: f32[]) -> f32[] {
  %Arg_0.1731 = f32[] parameter(0)
  %Arg_1.1732 = f32[] parameter(1)
  ROOT %add.1733 = f32[] add(f32[] %Arg_0.1731, f32[] %Arg_1.1732), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.10 (param_0.659: f32[3,48,4096,33], param_1.870: f32[3,3,4096,33]) -> f32[3,48,128,128] {
  %constant_514 = f32[] constant(0)
  %broadcast.971 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_514), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 48, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.163 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1409 = s32[] constant(0)
  %broadcast.1827 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1409), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.870 = f32[3,3,4096,33]{3,2,0,1} parameter(1)
  %slice.500 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.870), slice={[0:3], [1:2], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 1, 4096, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.356 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.500), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.259 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.356), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.154 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.259), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1408 = s32[] constant(1)
  %broadcast.1826 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1408), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.805 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.154, s32[3,4096,33]{2,1,0} %broadcast.1826), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1407 = s32[] constant(127)
  %broadcast.1825 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1407), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.143 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1827, s32[3,4096,33]{2,1,0} %add.805, s32[3,4096,33]{2,1,0} %broadcast.1825), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.269 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.143, s32[3,4096,33]{2,1,0} %broadcast.1827), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1406 = s32[] constant(128)
  %broadcast.1824 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1406), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.804 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.143, s32[3,4096,33]{2,1,0} %broadcast.1824), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.317 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.269, s32[3,4096,33]{2,1,0} %add.804, s32[3,4096,33]{2,1,0} %clamp.143), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.355 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.317), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.499 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.870), slice={[0:3], [2:3], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 1, 0, 0) limit_indices=(3, 2, 4096, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.354 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.499), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.258 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.354), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.153 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.258), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.803 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.153, s32[3,4096,33]{2,1,0} %broadcast.1826), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.142 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1827, s32[3,4096,33]{2,1,0} %add.803, s32[3,4096,33]{2,1,0} %broadcast.1825), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.268 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.142, s32[3,4096,33]{2,1,0} %broadcast.1827), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.802 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.142, s32[3,4096,33]{2,1,0} %broadcast.1824), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.316 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.268, s32[3,4096,33]{2,1,0} %add.802, s32[3,4096,33]{2,1,0} %clamp.142), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.353 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.316), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.214 = s32[3,4096,33,3]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.163, s32[3,4096,33,1]{3,2,1,0} %bitcast.355, s32[3,4096,33,1]{3,2,1,0} %bitcast.353), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.268 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.356, f32[3,4096,33]{2,1,0} %floor.259), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.267 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.354, f32[3,4096,33]{2,1,0} %floor.258), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.520 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.268, f32[3,4096,33]{2,1,0} %subtract.267), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1749 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.520), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.659 = f32[3,48,4096,33]{3,2,1,0} parameter(0)
  %multiply.281 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.1749, f32[3,48,4096,33]{3,2,1,0} %param_0.659), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.52 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.971, s32[3,4096,33,3]{3,2,1,0} %concatenate.214, f32[3,48,4096,33]{3,2,1,0} %multiply.281), update_window_dims={1}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_32.1730, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%fused_computation.30 (param_0.703: f32[3,48,128], param_1.923: f32[3,3,4096,33], param_2.934: f32[135168,144]) -> f32[3,48,4096,33] {
  %constant_1312 = f32[] constant(1)
  %broadcast.1725 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1312), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.923 = f32[3,3,4096,33]{3,2,0,1} parameter(1)
  %slice.442 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.923), slice={[0:3], [0:1], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1, 2, 3), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(3, 1, 4096, 33) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=165}
  %bitcast.284 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.442), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.227 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.284), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.254 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.284, f32[3,4096,33]{2,1,0} %floor.227), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.253 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1725, f32[3,4096,33]{2,1,0} %subtract.254), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1724 = f32[3,48,4096,33]{2,1,0,3} broadcast(f32[3,4096,33]{2,1,0} %subtract.253), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.703 = f32[3,48,128]{2,1,0} parameter(0)
  %iota.155 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1308 = s32[] constant(0)
  %broadcast.1723 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1308), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.134 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.227), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1306 = s32[] constant(127)
  %broadcast.1722 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1306), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.123 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1723, s32[3,4096,33]{2,1,0} %convert.134, s32[3,4096,33]{2,1,0} %broadcast.1722), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.249 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.123, s32[3,4096,33]{2,1,0} %broadcast.1723), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1305 = s32[] constant(128)
  %broadcast.1721 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1305), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.761 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.123, s32[3,4096,33]{2,1,0} %broadcast.1721), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.296 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.249, s32[3,4096,33]{2,1,0} %add.761, s32[3,4096,33]{2,1,0} %clamp.123), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.282 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.296), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.206 = s32[3,4096,33,2]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.155, s32[3,4096,33,1]{2,1,0,3} %bitcast.282), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.99 = f32[3,48,4096,33]{2,1,0,3} gather(f32[3,48,128]{2,1,0} %param_0.703, s32[3,4096,33,2]{2,1,0,3} %concatenate.206), offset_dims={1}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,48,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 48, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.512 = f32[3,48,4096,33]{2,1,0,3} multiply(f32[3,48,4096,33]{2,1,0,3} %broadcast.1724, f32[3,48,4096,33]{2,1,0,3} %gather.99), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.31 = f32[3,48,4096,33]{3,2,1,0} copy(f32[3,48,4096,33]{2,1,0,3} %multiply.512), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1720 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %subtract.254), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.154 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1307 = s32[] constant(1)
  %broadcast.1718 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1307), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.760 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.134, s32[3,4096,33]{2,1,0} %broadcast.1718), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.122 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1723, s32[3,4096,33]{2,1,0} %add.760, s32[3,4096,33]{2,1,0} %broadcast.1722), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.248 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.122, s32[3,4096,33]{2,1,0} %broadcast.1723), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.759 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.122, s32[3,4096,33]{2,1,0} %broadcast.1721), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.295 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.248, s32[3,4096,33]{2,1,0} %add.759, s32[3,4096,33]{2,1,0} %clamp.122), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.279 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.295), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.205 = s32[3,4096,33,2]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.154, s32[3,4096,33,1]{3,2,1,0} %bitcast.279), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.98 = f32[3,48,4096,33]{3,2,1,0} gather(f32[3,48,128]{2,1,0} %param_0.703, s32[3,4096,33,2]{3,2,1,0} %concatenate.205), offset_dims={1}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,48,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 48, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.511 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.1720, f32[3,48,4096,33]{3,2,1,0} %gather.98), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.758 = f32[3,48,4096,33]{3,2,1,0} add(f32[3,48,4096,33]{3,2,1,0} %copy.31, f32[3,48,4096,33]{3,2,1,0} %multiply.511), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.934 = f32[135168,144]{1,0} parameter(2)
  %reshape.592 = f32[4096,33,144]{0,2,1} reshape(f32[135168,144]{1,0} %param_2.934), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reshape[new_sizes=(4096, 33, 144) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
  %transpose.16 = f32[144,4096,33]{1,0,2} transpose(f32[4096,33,144]{0,2,1} %reshape.592), dimensions={2,0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/transpose[permutation=(2, 0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
  %bitcast.548 = f32[3,48,4096,33]{2,1,0,3} bitcast(f32[144,4096,33]{1,0,2} %transpose.16), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 48, 4096, 33) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %copy.25 = f32[3,48,4096,33]{3,2,1,0} copy(f32[3,48,4096,33]{2,1,0,3} %bitcast.548), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 48, 4096, 33) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  ROOT %multiply.282 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %add.758, f32[3,48,4096,33]{3,2,1,0} %copy.25), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
}

%fused_computation.31 (param_0.79: f32[3,48,128], param_1.979: f32[], param_2.1002: f32[], param_3.887: f32[], param_4.661: f32[3,48,128], param_5.622: f32[3,48,128], param_6.560: f32[3,48,128], param_7.512: f32[3,48,128]) -> (f32[3,48,128], f32[3,48,128], f32[3,48,128]) {
  %param_0.79 = f32[3,48,128]{2,1,0} parameter(0)
  %param_5.622 = f32[3,48,128]{2,1,0} parameter(5)
  %param_6.560 = f32[3,48,128]{2,1,0} parameter(6)
  %add.497.clone.1 = f32[3,48,128]{2,1,0} add(f32[3,48,128]{2,1,0} %param_5.622, f32[3,48,128]{2,1,0} %param_6.560), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_520_clone_1 = f32[] constant(0.0100000007)
  %broadcast.838.clone.1 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_520_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.290.clone.1 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %add.497.clone.1, f32[3,48,128]{2,1,0} %broadcast.838.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_7.512 = f32[3,48,128]{2,1,0} parameter(7)
  %constant_521_clone_1 = f32[] constant(0.9)
  %broadcast.837.clone.1 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_521_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.289.clone.1 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %param_7.512, f32[3,48,128]{2,1,0} %broadcast.837.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.376.clone.1 = f32[3,48,128]{2,1,0} add(f32[3,48,128]{2,1,0} %multiply.290.clone.1, f32[3,48,128]{2,1,0} %multiply.289.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.1002 = f32[] parameter(2)
  %broadcast.832 = f32[3,48,128]{2,1,0} broadcast(f32[] %param_2.1002), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_519_clone_1 = f32[] constant(0.1)
  %broadcast.836.clone.1 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_519_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.288.clone.1 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %add.497.clone.1, f32[3,48,128]{2,1,0} %broadcast.836.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.287.clone.1 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %multiply.288.clone.1, f32[3,48,128]{2,1,0} %multiply.288.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_517_clone_1 = f32[] constant(0.01)
  %broadcast.834.clone.1 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_517_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.286.clone.1 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %multiply.287.clone.1, f32[3,48,128]{2,1,0} %broadcast.834.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.661 = f32[3,48,128]{2,1,0} parameter(4)
  %constant_518_clone_1 = f32[] constant(0.99)
  %broadcast.833.clone.1 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_518_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.285.clone.1 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %param_4.661, f32[3,48,128]{2,1,0} %broadcast.833.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.375.clone.1 = f32[3,48,128]{2,1,0} add(f32[3,48,128]{2,1,0} %multiply.286.clone.1, f32[3,48,128]{2,1,0} %multiply.285.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.979 = f32[] parameter(1)
  %broadcast.831 = f32[3,48,128]{2,1,0} broadcast(f32[] %param_1.979), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.32 = f32[3,48,128]{2,1,0} divide(f32[3,48,128]{2,1,0} %add.375.clone.1, f32[3,48,128]{2,1,0} %broadcast.831), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.3 = f32[3,48,128]{2,1,0} sqrt(f32[3,48,128]{2,1,0} %divide.32), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_516 = f32[] constant(1e-08)
  %broadcast.830 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_516), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.373 = f32[3,48,128]{2,1,0} add(f32[3,48,128]{2,1,0} %sqrt.3, f32[3,48,128]{2,1,0} %broadcast.830), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.284 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %broadcast.832, f32[3,48,128]{2,1,0} %add.373)
  %divide.31 = f32[3,48,128]{2,1,0} divide(f32[3,48,128]{2,1,0} %add.376.clone.1, f32[3,48,128]{2,1,0} %multiply.284), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.887 = f32[] parameter(3)
  %constant_515 = f32[] constant(-0.02)
  %multiply.400 = f32[] multiply(f32[] %param_3.887, f32[] %constant_515), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.829 = f32[3,48,128]{2,1,0} broadcast(f32[] %multiply.400), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.283 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %divide.31, f32[3,48,128]{2,1,0} %broadcast.829), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.372 = f32[3,48,128]{2,1,0} add(f32[3,48,128]{2,1,0} %param_0.79, f32[3,48,128]{2,1,0} %multiply.283), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.54 = (f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}) tuple(f32[3,48,128]{2,1,0} %add.372, f32[3,48,128]{2,1,0} %add.375.clone.1, f32[3,48,128]{2,1,0} %add.376.clone.1)
}

%region_37.1785 (Arg_0.1786: f32[], Arg_1.1787: f32[]) -> f32[] {
  %Arg_0.1786 = f32[] parameter(0)
  %Arg_1.1787 = f32[] parameter(1)
  ROOT %add.1788 = f32[] add(f32[] %Arg_0.1786, f32[] %Arg_1.1787), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.11 (param_0.614: f32[3,48,4096,33], param_1.830: f32[3,3,4096,33]) -> f32[3,48,128] {
  %constant_522 = f32[] constant(0)
  %broadcast.972 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_522), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 48, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.151 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1295 = s32[] constant(0)
  %broadcast.1705 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1295), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.830 = f32[3,3,4096,33]{3,2,0,1} parameter(1)
  %slice.434 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.830), slice={[0:3], [0:1], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1, 2, 3), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(3, 1, 4096, 33) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=165}
  %bitcast.272 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.434), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.219 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.272), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.130 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.219), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1294 = s32[] constant(127)
  %broadcast.1704 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1294), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.119 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1705, s32[3,4096,33]{2,1,0} %convert.130, s32[3,4096,33]{2,1,0} %broadcast.1704), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.245 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.119, s32[3,4096,33]{2,1,0} %broadcast.1705), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1293 = s32[] constant(128)
  %broadcast.1703 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1293), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.753 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.119, s32[3,4096,33]{2,1,0} %broadcast.1703), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.292 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.245, s32[3,4096,33]{2,1,0} %add.753, s32[3,4096,33]{2,1,0} %clamp.119), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.271 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.292), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.202 = s32[3,4096,33,2]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.151, s32[3,4096,33,1]{2,1,0,3} %bitcast.271), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1281 = f32[] constant(1)
  %broadcast.1693 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1281), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.248 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.272, f32[3,4096,33]{2,1,0} %floor.219), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.247 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1693, f32[3,4096,33]{2,1,0} %subtract.248), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1692 = f32[3,48,4096,33]{2,1,0,3} broadcast(f32[3,4096,33]{2,1,0} %subtract.247), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.614 = f32[3,48,4096,33]{2,1,0,3} parameter(0)
  %multiply.291 = f32[3,48,4096,33]{2,1,0,3} multiply(f32[3,48,4096,33]{2,1,0,3} %broadcast.1692, f32[3,48,4096,33]{2,1,0,3} %param_0.614), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.53 = f32[3,48,128]{2,1,0} scatter(f32[3,48,128]{2,1,0} %broadcast.972, s32[3,4096,33,2]{2,1,0,3} %concatenate.202, f32[3,48,4096,33]{2,1,0,3} %multiply.291), update_window_dims={1}, inserted_window_dims={0,2}, scatter_dims_to_operand_dims={0,2}, index_vector_dim=3, to_apply=%region_37.1785, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_36.1781 (Arg_0.1782: f32[], Arg_1.1783: f32[]) -> f32[] {
  %Arg_0.1782 = f32[] parameter(0)
  %Arg_1.1783 = f32[] parameter(1)
  ROOT %add.1784 = f32[] add(f32[] %Arg_0.1782, f32[] %Arg_1.1783), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.12 (param_0.609: f32[3,48,4096,33], param_1.826: f32[3,3,4096,33]) -> f32[3,48,128] {
  %constant_523 = f32[] constant(0)
  %broadcast.973 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_523), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 48, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.147 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1275 = s32[] constant(0)
  %broadcast.1685 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1275), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.826 = f32[3,3,4096,33]{3,2,0,1} parameter(1)
  %slice.426 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.826), slice={[0:3], [0:1], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1, 2, 3), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(3, 1, 4096, 33) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=165}
  %bitcast.260 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.426), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.211 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.260), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.126 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.211), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1274 = s32[] constant(1)
  %broadcast.1684 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1274), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.749 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.126, s32[3,4096,33]{2,1,0} %broadcast.1684), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1273 = s32[] constant(127)
  %broadcast.1683 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1273), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.115 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1685, s32[3,4096,33]{2,1,0} %add.749, s32[3,4096,33]{2,1,0} %broadcast.1683), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.241 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.115, s32[3,4096,33]{2,1,0} %broadcast.1685), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1272 = s32[] constant(128)
  %broadcast.1682 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1272), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.748 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.115, s32[3,4096,33]{2,1,0} %broadcast.1682), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.287 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.241, s32[3,4096,33]{2,1,0} %add.748, s32[3,4096,33]{2,1,0} %clamp.115), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.259 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.287), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.198 = s32[3,4096,33,2]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.147, s32[3,4096,33,1]{3,2,1,0} %bitcast.259), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.78 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.260, f32[3,4096,33]{2,1,0} %floor.211), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.974 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %subtract.78), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.609 = f32[3,48,4096,33]{2,1,0,3} parameter(0)
  %copy.26 = f32[3,48,4096,33]{3,2,1,0} copy(f32[3,48,4096,33]{2,1,0,3} %param_0.609), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %multiply.292 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.974, f32[3,48,4096,33]{3,2,1,0} %copy.26), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.54 = f32[3,48,128]{2,1,0} scatter(f32[3,48,128]{2,1,0} %broadcast.973, s32[3,4096,33,2]{3,2,1,0} %concatenate.198, f32[3,48,4096,33]{3,2,1,0} %multiply.292), update_window_dims={1}, inserted_window_dims={0,2}, scatter_dims_to_operand_dims={0,2}, index_vector_dim=3, to_apply=%region_36.1781, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%fused_computation.34 (param_0.701: f32[3,48,128,128], param_1.922: f32[3,3,4096,33], param_2.931: f32[135168,144]) -> f32[3,48,4096,33] {
  %param_2.931 = f32[135168,144]{1,0} parameter(2)
  %reshape.590 = f32[4096,33,144]{0,2,1} reshape(f32[135168,144]{1,0} %param_2.931), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reshape[new_sizes=(4096, 33, 144) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
  %transpose.14 = f32[144,4096,33]{1,0,2} transpose(f32[4096,33,144]{0,2,1} %reshape.590), dimensions={2,0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/transpose[permutation=(2, 0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
  %bitcast.546 = f32[3,48,4096,33]{2,1,0,3} bitcast(f32[144,4096,33]{1,0,2} %transpose.14), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 48, 4096, 33) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %constant_1594 = f32[] constant(1)
  %broadcast.2041 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1594), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.922 = f32[3,3,4096,33]{3,2,0,1} parameter(1)
  %slice.618 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.922), slice={[0:3], [1:2], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 1, 4096, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.488 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.618), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.348 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.488), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.334 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.488, f32[3,4096,33]{2,1,0} %floor.348), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.333 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.2041, f32[3,4096,33]{2,1,0} %subtract.334), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.616 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.922), slice={[0:3], [2:3], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 1, 0, 0) limit_indices=(3, 2, 4096, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.487 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.616), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.347 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.487), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.332 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.487, f32[3,4096,33]{2,1,0} %floor.347), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.331 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.2041, f32[3,4096,33]{2,1,0} %subtract.332), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.549 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.333, f32[3,4096,33]{2,1,0} %subtract.331), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2040 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.549), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.701 = f32[3,48,128,128]{3,2,1,0} parameter(0)
  %iota.183 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1582 = s32[] constant(0)
  %broadcast.2039 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1582), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.202 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.348), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1580 = s32[] constant(127)
  %broadcast.2038 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1580), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.191 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.2039, s32[3,4096,33]{2,1,0} %convert.202, s32[3,4096,33]{2,1,0} %broadcast.2038), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.319 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.191, s32[3,4096,33]{2,1,0} %broadcast.2039), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1579 = s32[] constant(128)
  %broadcast.2037 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1579), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.885 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.191, s32[3,4096,33]{2,1,0} %broadcast.2037), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.365 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.319, s32[3,4096,33]{2,1,0} %add.885, s32[3,4096,33]{2,1,0} %clamp.191), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.485 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.365), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.201 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.347), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.190 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.2039, s32[3,4096,33]{2,1,0} %convert.201, s32[3,4096,33]{2,1,0} %broadcast.2038), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.318 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.190, s32[3,4096,33]{2,1,0} %broadcast.2039), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.884 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.190, s32[3,4096,33]{2,1,0} %broadcast.2037), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.364 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.318, s32[3,4096,33]{2,1,0} %add.884, s32[3,4096,33]{2,1,0} %clamp.190), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.483 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.364), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.235 = s32[3,4096,33,3]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.183, s32[3,4096,33,1]{2,1,0,3} %bitcast.485, s32[3,4096,33,1]{2,1,0,3} %bitcast.483), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.111 = f32[3,48,4096,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_0.701, s32[3,4096,33,3]{2,1,0,3} %concatenate.235), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.548 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.2040, f32[3,48,4096,33]{3,2,1,0} %gather.111), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.547 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.333, f32[3,4096,33]{2,1,0} %subtract.332), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2032 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.547), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.182 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.479 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.365), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1581 = s32[] constant(1)
  %broadcast.2027 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1581), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.881 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.201, s32[3,4096,33]{2,1,0} %broadcast.2027), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.188 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.2039, s32[3,4096,33]{2,1,0} %add.881, s32[3,4096,33]{2,1,0} %broadcast.2038), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.316 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.188, s32[3,4096,33]{2,1,0} %broadcast.2039), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.880 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.188, s32[3,4096,33]{2,1,0} %broadcast.2037), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.362 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.316, s32[3,4096,33]{2,1,0} %add.880, s32[3,4096,33]{2,1,0} %clamp.188), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.477 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.362), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.234 = s32[3,4096,33,3]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.182, s32[3,4096,33,1]{3,2,1,0} %bitcast.479, s32[3,4096,33,1]{3,2,1,0} %bitcast.477), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.110 = f32[3,48,4096,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_0.701, s32[3,4096,33,3]{3,2,1,0} %concatenate.234), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.546 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.2032, f32[3,48,4096,33]{3,2,1,0} %gather.110), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.879 = f32[3,48,4096,33]{3,2,1,0} add(f32[3,48,4096,33]{3,2,1,0} %multiply.548, f32[3,48,4096,33]{3,2,1,0} %multiply.546), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.544 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.334, f32[3,4096,33]{2,1,0} %subtract.331), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2023 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.544), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.181 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.878 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.202, s32[3,4096,33]{2,1,0} %broadcast.2027), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.187 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.2039, s32[3,4096,33]{2,1,0} %add.878, s32[3,4096,33]{2,1,0} %broadcast.2038), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.315 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.187, s32[3,4096,33]{2,1,0} %broadcast.2039), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.877 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.187, s32[3,4096,33]{2,1,0} %broadcast.2037), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.361 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.315, s32[3,4096,33]{2,1,0} %add.877, s32[3,4096,33]{2,1,0} %clamp.187), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.473 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.361), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.233 = s32[3,4096,33,3]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.181, s32[3,4096,33,1]{2,1,0,3} %bitcast.473, s32[3,4096,33,1]{2,1,0,3} %bitcast.483), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.109 = f32[3,48,4096,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_0.701, s32[3,4096,33,3]{2,1,0,3} %concatenate.233), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.543 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.2023, f32[3,48,4096,33]{3,2,1,0} %gather.109), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.875 = f32[3,48,4096,33]{3,2,1,0} add(f32[3,48,4096,33]{3,2,1,0} %add.879, f32[3,48,4096,33]{3,2,1,0} %multiply.543), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.542 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.334, f32[3,4096,33]{2,1,0} %subtract.332), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2015 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.542), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.180 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.467 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.361), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.232 = s32[3,4096,33,3]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.180, s32[3,4096,33,1]{3,2,1,0} %bitcast.467, s32[3,4096,33,1]{3,2,1,0} %bitcast.477), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.108 = f32[3,48,4096,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_0.701, s32[3,4096,33,3]{3,2,1,0} %concatenate.232), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.541 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.2015, f32[3,48,4096,33]{3,2,1,0} %gather.108), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.870 = f32[3,48,4096,33]{3,2,1,0} add(f32[3,48,4096,33]{3,2,1,0} %add.875, f32[3,48,4096,33]{3,2,1,0} %multiply.541), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.27 = f32[3,48,4096,33]{2,1,0,3} copy(f32[3,48,4096,33]{3,2,1,0} %add.870), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %multiply.293 = f32[3,48,4096,33]{2,1,0,3} multiply(f32[3,48,4096,33]{2,1,0,3} %bitcast.546, f32[3,48,4096,33]{2,1,0,3} %copy.27), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
}

%fused_computation.36 (param_0.90: f32[128,3], param_1.980: f32[], param_2.1004: f32[], param_3.889: f32[], param_4.663: f32[128,3], param_5.625: f32[128,3], param_6.567: f32[128,3]) -> (f32[128,3], f32[128,3], f32[128,3]) {
  %param_0.90 = f32[128,3]{1,0} parameter(0)
  %param_5.625 = f32[128,3]{1,0} parameter(5)
  %constant_530_clone_1 = f32[] constant(0.0100000007)
  %broadcast.849.clone.1 = f32[128,3]{1,0} broadcast(f32[] %constant_530_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.302.clone.1 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %param_5.625, f32[128,3]{1,0} %broadcast.849.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_6.567 = f32[128,3]{1,0} parameter(6)
  %constant_533_clone_1 = f32[] constant(0.9)
  %broadcast.848.clone.1 = f32[128,3]{1,0} broadcast(f32[] %constant_533_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.301.clone.1 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %param_6.567, f32[128,3]{1,0} %broadcast.848.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.382.clone.1 = f32[128,3]{1,0} add(f32[128,3]{1,0} %multiply.302.clone.1, f32[128,3]{1,0} %multiply.301.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.1004 = f32[] parameter(2)
  %broadcast.844 = f32[128,3]{1,0} broadcast(f32[] %param_2.1004), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_528_clone_1 = f32[] constant(0.1)
  %broadcast.847.clone.1 = f32[128,3]{1,0} broadcast(f32[] %constant_528_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.300.clone.1 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %param_5.625, f32[128,3]{1,0} %broadcast.847.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.299.clone.1 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %multiply.300.clone.1, f32[128,3]{1,0} %multiply.300.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_526_clone_1 = f32[] constant(0.01)
  %broadcast.846.clone.1 = f32[128,3]{1,0} broadcast(f32[] %constant_526_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.297.clone.1 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %multiply.299.clone.1, f32[128,3]{1,0} %broadcast.846.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.663 = f32[128,3]{1,0} parameter(4)
  %constant_527_clone_1 = f32[] constant(0.99)
  %broadcast.845.clone.1 = f32[128,3]{1,0} broadcast(f32[] %constant_527_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.296.clone.1 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %param_4.663, f32[128,3]{1,0} %broadcast.845.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.380.clone.1 = f32[128,3]{1,0} add(f32[128,3]{1,0} %multiply.297.clone.1, f32[128,3]{1,0} %multiply.296.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.980 = f32[] parameter(1)
  %broadcast.843 = f32[128,3]{1,0} broadcast(f32[] %param_1.980), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.34 = f32[128,3]{1,0} divide(f32[128,3]{1,0} %add.380.clone.1, f32[128,3]{1,0} %broadcast.843), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.4 = f32[128,3]{1,0} sqrt(f32[128,3]{1,0} %divide.34), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_524 = f32[] constant(1e-08)
  %broadcast.841 = f32[128,3]{1,0} broadcast(f32[] %constant_524), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.379 = f32[128,3]{1,0} add(f32[128,3]{1,0} %sqrt.4, f32[128,3]{1,0} %broadcast.841), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.295 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %broadcast.844, f32[128,3]{1,0} %add.379)
  %divide.33 = f32[128,3]{1,0} divide(f32[128,3]{1,0} %add.382.clone.1, f32[128,3]{1,0} %multiply.295), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.889 = f32[] parameter(3)
  %constant_525 = f32[] constant(-0.001)
  %multiply.401 = f32[] multiply(f32[] %param_3.889, f32[] %constant_525), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.839 = f32[128,3]{1,0} broadcast(f32[] %multiply.401), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.294 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %divide.33, f32[128,3]{1,0} %broadcast.839), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.377 = f32[128,3]{1,0} add(f32[128,3]{1,0} %param_0.90, f32[128,3]{1,0} %multiply.294), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.56 = (f32[128,3]{1,0}, f32[128,3]{1,0}, f32[128,3]{1,0}) tuple(f32[128,3]{1,0} %add.377, f32[128,3]{1,0} %add.380.clone.1, f32[128,3]{1,0} %add.382.clone.1)
}

%fused_computation.39 (param_0.95: f32[3], param_1.981: f32[], param_2.1006: f32[], param_3.891: f32[], param_4.665: f32[3], param_5.628: f32[3], param_6.574: f32[3]) -> (f32[3], f32[3], f32[3]) {
  %param_0.95 = f32[3]{0} parameter(0)
  %param_5.628 = f32[3]{0} parameter(5)
  %constant_539_clone_1 = f32[] constant(0.0100000007)
  %broadcast.861.clone.1 = f32[3]{0} broadcast(f32[] %constant_539_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.310.clone.1 = f32[3]{0} multiply(f32[3]{0} %param_5.628, f32[3]{0} %broadcast.861.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_6.574 = f32[3]{0} parameter(6)
  %constant_540_clone_1 = f32[] constant(0.9)
  %broadcast.859.clone.1 = f32[3]{0} broadcast(f32[] %constant_540_clone_1), dimensions={}
  %multiply.309.clone.1 = f32[3]{0} multiply(f32[3]{0} %param_6.574, f32[3]{0} %broadcast.859.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.386.clone.1 = f32[3]{0} add(f32[3]{0} %multiply.310.clone.1, f32[3]{0} %multiply.309.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.1006 = f32[] parameter(2)
  %broadcast.854 = f32[3]{0} broadcast(f32[] %param_2.1006), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_538_clone_1 = f32[] constant(0.1)
  %broadcast.857.clone.1 = f32[3]{0} broadcast(f32[] %constant_538_clone_1), dimensions={}
  %multiply.308.clone.1 = f32[3]{0} multiply(f32[3]{0} %param_5.628, f32[3]{0} %broadcast.857.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.307.clone.1 = f32[3]{0} multiply(f32[3]{0} %multiply.308.clone.1, f32[3]{0} %multiply.308.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_536_clone_1 = f32[] constant(0.01)
  %broadcast.856.clone.1 = f32[3]{0} broadcast(f32[] %constant_536_clone_1), dimensions={}
  %multiply.306.clone.1 = f32[3]{0} multiply(f32[3]{0} %multiply.307.clone.1, f32[3]{0} %broadcast.856.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.665 = f32[3]{0} parameter(4)
  %constant_537_clone_1 = f32[] constant(0.99)
  %broadcast.855.clone.1 = f32[3]{0} broadcast(f32[] %constant_537_clone_1), dimensions={}
  %multiply.305.clone.1 = f32[3]{0} multiply(f32[3]{0} %param_4.665, f32[3]{0} %broadcast.855.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.385.clone.1 = f32[3]{0} add(f32[3]{0} %multiply.306.clone.1, f32[3]{0} %multiply.305.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.981 = f32[] parameter(1)
  %broadcast.852 = f32[3]{0} broadcast(f32[] %param_1.981), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.36 = f32[3]{0} divide(f32[3]{0} %add.385.clone.1, f32[3]{0} %broadcast.852), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.5 = f32[3]{0} sqrt(f32[3]{0} %divide.36), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_534 = f32[] constant(1e-08)
  %broadcast.851 = f32[3]{0} broadcast(f32[] %constant_534), dimensions={}
  %add.384 = f32[3]{0} add(f32[3]{0} %sqrt.5, f32[3]{0} %broadcast.851), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.304 = f32[3]{0} multiply(f32[3]{0} %broadcast.854, f32[3]{0} %add.384)
  %divide.35 = f32[3]{0} divide(f32[3]{0} %add.386.clone.1, f32[3]{0} %multiply.304), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.891 = f32[] parameter(3)
  %constant_535 = f32[] constant(-0.001)
  %multiply.402 = f32[] multiply(f32[] %param_3.891, f32[] %constant_535), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.850 = f32[3]{0} broadcast(f32[] %multiply.402), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.303 = f32[3]{0} multiply(f32[3]{0} %divide.35, f32[3]{0} %broadcast.850), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.383 = f32[3]{0} add(f32[3]{0} %param_0.95, f32[3]{0} %multiply.303), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.58 = (f32[3]{0}, f32[3]{0}, f32[3]{0}) tuple(f32[3]{0} %add.383, f32[3]{0} %add.385.clone.1, f32[3]{0} %add.386.clone.1)
}

%input_fused_computation_reduce.1 (param_0.384: f32[135168,3]) -> f32[3,368] {
  %param_0.384 = f32[135168,3]{0,1} parameter(0)
  %bitcast.82 = f32[3,135168]{1,0} bitcast(f32[135168,3]{0,1} %param_0.384), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(sigmoid)))/jit(jit_transpose(jvp(sigmoid)))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %constant_541 = f32[] constant(0)
  %pad.110 = f32[3,135424]{1,0} pad(f32[3,135168]{1,0} %bitcast.82, f32[] %constant_541), padding=0_0x0_256
  %bitcast.81 = f32[3,368,368]{2,1,0} bitcast(f32[3,135424]{1,0} %pad.110)
  ROOT %reduce.40 = f32[3,368]{1,0} reduce(f32[3,368,368]{2,1,0} %bitcast.81, f32[] %constant_541), dimensions={2}, to_apply=%region_27.1710
}

%fused_computation.42 (param_0.102: f32[128,128], param_1.982: f32[], param_2.1008: f32[], param_3.893: f32[], param_4.667: f32[128,128], param_5.631: f32[128,128], param_6.581: f32[128,128]) -> (f32[128,128], f32[128,128], f32[128,128]) {
  %param_0.102 = f32[128,128]{1,0} parameter(0)
  %param_5.631 = f32[128,128]{1,0} parameter(5)
  %constant_547_clone_1 = f32[] constant(0.0100000007)
  %broadcast.870.clone.1 = f32[128,128]{1,0} broadcast(f32[] %constant_547_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.319.clone.1 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %param_5.631, f32[128,128]{1,0} %broadcast.870.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_6.581 = f32[128,128]{1,0} parameter(6)
  %constant_548_clone_1 = f32[] constant(0.9)
  %broadcast.869.clone.1 = f32[128,128]{1,0} broadcast(f32[] %constant_548_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.318.clone.1 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %param_6.581, f32[128,128]{1,0} %broadcast.869.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.390.clone.1 = f32[128,128]{1,0} add(f32[128,128]{1,0} %multiply.319.clone.1, f32[128,128]{1,0} %multiply.318.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.1008 = f32[] parameter(2)
  %broadcast.865 = f32[128,128]{1,0} broadcast(f32[] %param_2.1008), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_546_clone_1 = f32[] constant(0.1)
  %broadcast.868.clone.1 = f32[128,128]{1,0} broadcast(f32[] %constant_546_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.317.clone.1 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %param_5.631, f32[128,128]{1,0} %broadcast.868.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.316.clone.1 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %multiply.317.clone.1, f32[128,128]{1,0} %multiply.317.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_544_clone_1 = f32[] constant(0.01)
  %broadcast.867.clone.1 = f32[128,128]{1,0} broadcast(f32[] %constant_544_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.315.clone.1 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %multiply.316.clone.1, f32[128,128]{1,0} %broadcast.867.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.667 = f32[128,128]{1,0} parameter(4)
  %constant_545_clone_1 = f32[] constant(0.99)
  %broadcast.866.clone.1 = f32[128,128]{1,0} broadcast(f32[] %constant_545_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.314.clone.1 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %param_4.667, f32[128,128]{1,0} %broadcast.866.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.389.clone.1 = f32[128,128]{1,0} add(f32[128,128]{1,0} %multiply.315.clone.1, f32[128,128]{1,0} %multiply.314.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.982 = f32[] parameter(1)
  %broadcast.864 = f32[128,128]{1,0} broadcast(f32[] %param_1.982), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.38 = f32[128,128]{1,0} divide(f32[128,128]{1,0} %add.389.clone.1, f32[128,128]{1,0} %broadcast.864), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.6 = f32[128,128]{1,0} sqrt(f32[128,128]{1,0} %divide.38), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_542 = f32[] constant(1e-08)
  %broadcast.863 = f32[128,128]{1,0} broadcast(f32[] %constant_542), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.388 = f32[128,128]{1,0} add(f32[128,128]{1,0} %sqrt.6, f32[128,128]{1,0} %broadcast.863), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.313 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %broadcast.865, f32[128,128]{1,0} %add.388)
  %divide.37 = f32[128,128]{1,0} divide(f32[128,128]{1,0} %add.390.clone.1, f32[128,128]{1,0} %multiply.313), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.893 = f32[] parameter(3)
  %constant_543 = f32[] constant(-0.001)
  %multiply.403 = f32[] multiply(f32[] %param_3.893, f32[] %constant_543), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.862 = f32[128,128]{1,0} broadcast(f32[] %multiply.403), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.312 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %divide.37, f32[128,128]{1,0} %broadcast.862), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.387 = f32[128,128]{1,0} add(f32[128,128]{1,0} %param_0.102, f32[128,128]{1,0} %multiply.312), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.60 = (f32[128,128]{1,0}, f32[128,128]{1,0}, f32[128,128]{1,0}) tuple(f32[128,128]{1,0} %add.387, f32[128,128]{1,0} %add.389.clone.1, f32[128,128]{1,0} %add.390.clone.1)
}

%fused_computation.45 (param_0.107: f32[128], param_1.983: f32[], param_2.1010: f32[], param_3.895: f32[], param_4.669: f32[128], param_5.634: f32[128], param_6.588: f32[128], param_7.576: f32[128], param_8.417: f32[128], param_9.277: f32[128], param_10.186: f32[128]) -> (f32[128], f32[128], f32[128], f32[128], f32[128], /*index=5*/f32[128]) {
  %param_0.107 = f32[128]{0} parameter(0)
  %param_5.634 = f32[128]{0} parameter(5)
  %constant_555_clone_1 = f32[] constant(0.0100000007)
  %broadcast.985.clone.1 = f32[128]{0} broadcast(f32[] %constant_555_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.327.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_5.634, f32[128]{0} %broadcast.985.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_6.588 = f32[128]{0} parameter(6)
  %constant_556_clone_1 = f32[] constant(0.9)
  %broadcast.987.clone.1 = f32[128]{0} broadcast(f32[] %constant_556_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.326.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_6.588, f32[128]{0} %broadcast.987.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.394.clone.1 = f32[128]{0} add(f32[128]{0} %multiply.327.clone.1, f32[128]{0} %multiply.326.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.1010 = f32[] parameter(2)
  %broadcast.977 = f32[128]{0} broadcast(f32[] %param_2.1010), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_554_clone_1 = f32[] constant(0.1)
  %broadcast.980.clone.1 = f32[128]{0} broadcast(f32[] %constant_554_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.325.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_5.634, f32[128]{0} %broadcast.980.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.324.clone.1 = f32[128]{0} multiply(f32[128]{0} %multiply.325.clone.1, f32[128]{0} %multiply.325.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_552_clone_1 = f32[] constant(0.01)
  %broadcast.981.clone.1 = f32[128]{0} broadcast(f32[] %constant_552_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.323.clone.1 = f32[128]{0} multiply(f32[128]{0} %multiply.324.clone.1, f32[128]{0} %broadcast.981.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.669 = f32[128]{0} parameter(4)
  %constant_553_clone_1 = f32[] constant(0.99)
  %broadcast.983.clone.1 = f32[128]{0} broadcast(f32[] %constant_553_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.322.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_4.669, f32[128]{0} %broadcast.983.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.393.clone.1 = f32[128]{0} add(f32[128]{0} %multiply.323.clone.1, f32[128]{0} %multiply.322.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.983 = f32[] parameter(1)
  %broadcast.975 = f32[128]{0} broadcast(f32[] %param_1.983), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.40 = f32[128]{0} divide(f32[128]{0} %add.393.clone.1, f32[128]{0} %broadcast.975), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.7 = f32[128]{0} sqrt(f32[128]{0} %divide.40), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_550 = f32[] constant(1e-08)
  %broadcast.976 = f32[128]{0} broadcast(f32[] %constant_550), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.392 = f32[128]{0} add(f32[128]{0} %sqrt.7, f32[128]{0} %broadcast.976), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.321 = f32[128]{0} multiply(f32[128]{0} %broadcast.977, f32[128]{0} %add.392)
  %divide.39 = f32[128]{0} divide(f32[128]{0} %add.394.clone.1, f32[128]{0} %multiply.321), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.895 = f32[] parameter(3)
  %constant_551 = f32[] constant(-0.001)
  %multiply.404 = f32[] multiply(f32[] %param_3.895, f32[] %constant_551), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.979 = f32[128]{0} broadcast(f32[] %multiply.404), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.320 = f32[128]{0} multiply(f32[128]{0} %divide.39, f32[128]{0} %broadcast.979), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.391 = f32[128]{0} add(f32[128]{0} %param_0.107, f32[128]{0} %multiply.320), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %param_7.576 = f32[128]{0} parameter(7)
  %param_9.277 = f32[128]{0} parameter(9)
  %multiply.343.clone.1.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_9.277, f32[128]{0} %broadcast.985.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_10.186 = f32[128]{0} parameter(10)
  %multiply.342.clone.1.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_10.186, f32[128]{0} %broadcast.987.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.404.clone.1.clone.1 = f32[128]{0} add(f32[128]{0} %multiply.343.clone.1.clone.1, f32[128]{0} %multiply.342.clone.1.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.341.clone.1.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_9.277, f32[128]{0} %broadcast.980.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.340.clone.1.clone.1 = f32[128]{0} multiply(f32[128]{0} %multiply.341.clone.1.clone.1, f32[128]{0} %multiply.341.clone.1.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.339.clone.1.clone.1 = f32[128]{0} multiply(f32[128]{0} %multiply.340.clone.1.clone.1, f32[128]{0} %broadcast.981.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_8.417 = f32[128]{0} parameter(8)
  %multiply.338.clone.1.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_8.417, f32[128]{0} %broadcast.983.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.403.clone.1.clone.1 = f32[128]{0} add(f32[128]{0} %multiply.339.clone.1.clone.1, f32[128]{0} %multiply.338.clone.1.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %divide.44.clone.1 = f32[128]{0} divide(f32[128]{0} %add.403.clone.1.clone.1, f32[128]{0} %broadcast.975), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.9.clone.1 = f32[128]{0} sqrt(f32[128]{0} %divide.44.clone.1), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.402.clone.1 = f32[128]{0} add(f32[128]{0} %sqrt.9.clone.1, f32[128]{0} %broadcast.976), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.337.clone.1 = f32[128]{0} multiply(f32[128]{0} %broadcast.977, f32[128]{0} %add.402.clone.1)
  %divide.43.clone.1 = f32[128]{0} divide(f32[128]{0} %add.404.clone.1.clone.1, f32[128]{0} %multiply.337.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.336.clone.1 = f32[128]{0} multiply(f32[128]{0} %divide.43.clone.1, f32[128]{0} %broadcast.979), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.401.clone.1 = f32[128]{0} add(f32[128]{0} %param_7.576, f32[128]{0} %multiply.336.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.67 = (f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) tuple(f32[128]{0} %add.391, f32[128]{0} %add.393.clone.1, f32[128]{0} %add.394.clone.1, f32[128]{0} %add.401.clone.1, f32[128]{0} %add.403.clone.1.clone.1, /*index=5*/f32[128]{0} %add.404.clone.1.clone.1)
}

%region_28.1714 (Arg_0.1715: f32[], Arg_1.1716: f32[]) -> f32[] {
  %Arg_0.1715 = f32[] parameter(0)
  %Arg_1.1716 = f32[] parameter(1)
  ROOT %add.1717 = f32[] add(f32[] %Arg_0.1715, f32[] %Arg_1.1716), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
}

%region_29.1718 (Arg_0.1719: f32[], Arg_1.1720: f32[]) -> f32[] {
  %Arg_0.1719 = f32[] parameter(0)
  %Arg_1.1720 = f32[] parameter(1)
  ROOT %add.1721 = f32[] add(f32[] %Arg_0.1719, f32[] %Arg_1.1720), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
}

%input_fused_computation_reduce.2 (param_0.385: f32[135168,128], param_1.1018: f32[135168,128]) -> (f32[368,128], f32[368,128]) {
  %param_0.385 = f32[135168,128]{1,0} parameter(0)
  %constant_557 = f32[] constant(0)
  %pad.111 = f32[135424,128]{1,0} pad(f32[135168,128]{1,0} %param_0.385, f32[] %constant_557), padding=0_256x0_0
  %bitcast.83 = f32[368,368,128]{2,1,0} bitcast(f32[135424,128]{1,0} %pad.111)
  %reduce.41 = f32[368,128]{1,0} reduce(f32[368,368,128]{2,1,0} %bitcast.83, f32[] %constant_557), dimensions={0}, to_apply=%region_28.1714
  %param_1.1018 = f32[135168,128]{1,0} parameter(1)
  %pad.112.clone.1 = f32[135424,128]{1,0} pad(f32[135168,128]{1,0} %param_1.1018, f32[] %constant_557), padding=0_256x0_0
  %bitcast.84.clone.1 = f32[368,368,128]{2,1,0} bitcast(f32[135424,128]{1,0} %pad.112.clone.1)
  %reduce.42.clone.1 = f32[368,128]{1,0} reduce(f32[368,368,128]{2,1,0} %bitcast.84.clone.1, f32[] %constant_557), dimensions={0}, to_apply=%region_29.1718
  ROOT %tuple.90 = (f32[368,128]{1,0}, f32[368,128]{1,0}) tuple(f32[368,128]{1,0} %reduce.41, f32[368,128]{1,0} %reduce.42.clone.1)
}

%fused_computation.48 (param_0.114: f32[150,128], param_1.984: f32[], param_2.1012: f32[], param_3.897: f32[], param_4.671: f32[150,128], param_5.637: f32[150,128], param_6.595: f32[150,128]) -> (f32[150,128], f32[150,128], f32[150,128]) {
  %param_0.114 = f32[150,128]{1,0} parameter(0)
  %param_5.637 = f32[150,128]{1,0} parameter(5)
  %constant_563_clone_1 = f32[] constant(0.0100000007)
  %broadcast.884.clone.1 = f32[150,128]{1,0} broadcast(f32[] %constant_563_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.335.clone.1 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %param_5.637, f32[150,128]{1,0} %broadcast.884.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_6.595 = f32[150,128]{1,0} parameter(6)
  %constant_564_clone_1 = f32[] constant(0.9)
  %broadcast.883.clone.1 = f32[150,128]{1,0} broadcast(f32[] %constant_564_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.334.clone.1 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %param_6.595, f32[150,128]{1,0} %broadcast.883.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.400.clone.1 = f32[150,128]{1,0} add(f32[150,128]{1,0} %multiply.335.clone.1, f32[150,128]{1,0} %multiply.334.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.1012 = f32[] parameter(2)
  %broadcast.876 = f32[150,128]{1,0} broadcast(f32[] %param_2.1012), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_562_clone_1 = f32[] constant(0.1)
  %broadcast.882.clone.1 = f32[150,128]{1,0} broadcast(f32[] %constant_562_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.333.clone.1 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %param_5.637, f32[150,128]{1,0} %broadcast.882.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.332.clone.1 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %multiply.333.clone.1, f32[150,128]{1,0} %multiply.333.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_560_clone_1 = f32[] constant(0.01)
  %broadcast.880.clone.1 = f32[150,128]{1,0} broadcast(f32[] %constant_560_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.331.clone.1 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %multiply.332.clone.1, f32[150,128]{1,0} %broadcast.880.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.671 = f32[150,128]{1,0} parameter(4)
  %constant_561_clone_1 = f32[] constant(0.99)
  %broadcast.878.clone.1 = f32[150,128]{1,0} broadcast(f32[] %constant_561_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.330.clone.1 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %param_4.671, f32[150,128]{1,0} %broadcast.878.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.398.clone.1 = f32[150,128]{1,0} add(f32[150,128]{1,0} %multiply.331.clone.1, f32[150,128]{1,0} %multiply.330.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.984 = f32[] parameter(1)
  %broadcast.874 = f32[150,128]{1,0} broadcast(f32[] %param_1.984), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.42 = f32[150,128]{1,0} divide(f32[150,128]{1,0} %add.398.clone.1, f32[150,128]{1,0} %broadcast.874), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.8 = f32[150,128]{1,0} sqrt(f32[150,128]{1,0} %divide.42), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_558 = f32[] constant(1e-08)
  %broadcast.872 = f32[150,128]{1,0} broadcast(f32[] %constant_558), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.397 = f32[150,128]{1,0} add(f32[150,128]{1,0} %sqrt.8, f32[150,128]{1,0} %broadcast.872), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.329 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %broadcast.876, f32[150,128]{1,0} %add.397)
  %divide.41 = f32[150,128]{1,0} divide(f32[150,128]{1,0} %add.400.clone.1, f32[150,128]{1,0} %multiply.329), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.897 = f32[] parameter(3)
  %constant_559 = f32[] constant(-0.001)
  %multiply.405 = f32[] multiply(f32[] %param_3.897, f32[] %constant_559), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.871 = f32[150,128]{1,0} broadcast(f32[] %multiply.405), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.328 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %divide.41, f32[150,128]{1,0} %broadcast.871), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.395 = f32[150,128]{1,0} add(f32[150,128]{1,0} %param_0.114, f32[150,128]{1,0} %multiply.328), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.64 = (f32[150,128]{1,0}, f32[150,128]{1,0}, f32[150,128]{1,0}) tuple(f32[150,128]{1,0} %add.395, f32[150,128]{1,0} %add.398.clone.1, f32[150,128]{1,0} %add.400.clone.1)
}

%fused_computation.54 (param_0.126: f32[144,27], param_1.995: f32[], param_2.1026: f32[], param_3.905: f32[], param_4.679: f32[144,27], param_5.644: f32[144,27], param_6.609: f32[144,27]) -> (f32[144,27], f32[144,27], f32[144,27]) {
  %param_0.126 = f32[144,27]{1,0} parameter(0)
  %param_5.644 = f32[144,27]{1,0} parameter(5)
  %constant_583_clone_1 = f32[] constant(0.0100000007)
  %broadcast.893.clone.1 = f32[144,27]{1,0} broadcast(f32[] %constant_583_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.352.clone.1 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %param_5.644, f32[144,27]{1,0} %broadcast.893.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_6.609 = f32[144,27]{1,0} parameter(6)
  %constant_584_clone_1 = f32[] constant(0.9)
  %broadcast.892.clone.1 = f32[144,27]{1,0} broadcast(f32[] %constant_584_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.351.clone.1 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %param_6.609, f32[144,27]{1,0} %broadcast.892.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.409.clone.1 = f32[144,27]{1,0} add(f32[144,27]{1,0} %multiply.352.clone.1, f32[144,27]{1,0} %multiply.351.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.1026 = f32[] parameter(2)
  %broadcast.888 = f32[144,27]{1,0} broadcast(f32[] %param_2.1026), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_582_clone_1 = f32[] constant(0.1)
  %broadcast.891.clone.1 = f32[144,27]{1,0} broadcast(f32[] %constant_582_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.350.clone.1 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %param_5.644, f32[144,27]{1,0} %broadcast.891.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.349.clone.1 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %multiply.350.clone.1, f32[144,27]{1,0} %multiply.350.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_580_clone_1 = f32[] constant(0.01)
  %broadcast.890.clone.1 = f32[144,27]{1,0} broadcast(f32[] %constant_580_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.348.clone.1 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %multiply.349.clone.1, f32[144,27]{1,0} %broadcast.890.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.679 = f32[144,27]{1,0} parameter(4)
  %constant_581_clone_1 = f32[] constant(0.99)
  %broadcast.889.clone.1 = f32[144,27]{1,0} broadcast(f32[] %constant_581_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.347.clone.1 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %param_4.679, f32[144,27]{1,0} %broadcast.889.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.407.clone.1 = f32[144,27]{1,0} add(f32[144,27]{1,0} %multiply.348.clone.1, f32[144,27]{1,0} %multiply.347.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.995 = f32[] parameter(1)
  %broadcast.887 = f32[144,27]{1,0} broadcast(f32[] %param_1.995), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.46 = f32[144,27]{1,0} divide(f32[144,27]{1,0} %add.407.clone.1, f32[144,27]{1,0} %broadcast.887), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.10 = f32[144,27]{1,0} sqrt(f32[144,27]{1,0} %divide.46), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_573 = f32[] constant(1e-08)
  %broadcast.886 = f32[144,27]{1,0} broadcast(f32[] %constant_573), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.406 = f32[144,27]{1,0} add(f32[144,27]{1,0} %sqrt.10, f32[144,27]{1,0} %broadcast.886), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.345 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %broadcast.888, f32[144,27]{1,0} %add.406)
  %divide.45 = f32[144,27]{1,0} divide(f32[144,27]{1,0} %add.409.clone.1, f32[144,27]{1,0} %multiply.345), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.905 = f32[] parameter(3)
  %constant_574 = f32[] constant(-0.001)
  %multiply.408 = f32[] multiply(f32[] %param_3.905, f32[] %constant_574), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.885 = f32[144,27]{1,0} broadcast(f32[] %multiply.408), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.344 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %divide.45, f32[144,27]{1,0} %broadcast.885), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.405 = f32[144,27]{1,0} add(f32[144,27]{1,0} %param_0.126, f32[144,27]{1,0} %multiply.344), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.73 = (f32[144,27]{1,0}, f32[144,27]{1,0}, f32[144,27]{1,0}) tuple(f32[144,27]{1,0} %add.405, f32[144,27]{1,0} %add.407.clone.1, f32[144,27]{1,0} %add.409.clone.1)
}

%fused_computation.56 (param_0.746: s32[]) -> (f32[], f32[], s32[]) {
  %constant_578 = f32[] constant(1)
  %constant_577 = f32[] constant(0.99)
  %param_0.746 = s32[] parameter(0)
  %constant_456_clone_1 = s32[] constant(2147483647)
  %compare.132.clone.1 = pred[] compare(s32[] %param_0.746, s32[] %constant_456_clone_1), direction=LT, metadata={op_name="jit(training_step)/jit(main)/lt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/numerics.py" source_line=68}
  %constant_579_clone_1 = s32[] constant(1)
  %add.408.clone.1 = s32[] add(s32[] %param_0.746, s32[] %constant_579_clone_1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/numerics.py" source_line=68}
  %select.125.clone.1 = s32[] select(pred[] %compare.132.clone.1, s32[] %add.408.clone.1, s32[] %constant_456_clone_1), metadata={op_name="jit(training_step)/jit(main)/jit(_where)/jit(jit__where)/select_n" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/numerics.py" source_line=68}
  %convert.22 = f32[] convert(s32[] %select.125.clone.1), metadata={op_name="jit(training_step)/jit(main)/convert_element_type[new_dtype=float32 weak_type=True]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=86}
  %power.3 = f32[] power(f32[] %constant_577, f32[] %convert.22), metadata={op_name="jit(training_step)/jit(main)/pow" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=86}
  %subtract.64 = f32[] subtract(f32[] %constant_578, f32[] %power.3), metadata={op_name="jit(training_step)/jit(main)/sub" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=86}
  %constant_575_clone_1 = f32[] constant(0.9)
  %power.4.clone.1 = f32[] power(f32[] %constant_575_clone_1, f32[] %convert.22), metadata={op_name="jit(training_step)/jit(main)/pow" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=86}
  %subtract.65.clone.1 = f32[] subtract(f32[] %constant_578, f32[] %power.4.clone.1), metadata={op_name="jit(training_step)/jit(main)/sub" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=86}
  ROOT %tuple.71 = (f32[], f32[], s32[]) tuple(f32[] %subtract.64, f32[] %subtract.65.clone.1, s32[] %select.125.clone.1)
}

%region_30.1722 (Arg_0.1723: f32[], Arg_1.1724: f32[]) -> f32[] {
  %Arg_0.1723 = f32[] parameter(0)
  %Arg_1.1724 = f32[] parameter(1)
  ROOT %add.1725 = f32[] add(f32[] %Arg_0.1723, f32[] %Arg_1.1724), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
}

%fused_computation.61 (param_0.389: f32[135168,27,4], param_1.522: f32[135168,150], param_2.463: f32[2]) -> f32[135168,27] {
  %param_1.522 = f32[135168,150]{1,0} parameter(1)
  %slice.309 = f32[135168,27]{1,0} slice(f32[135168,150]{1,0} %param_1.522), slice={[0:135168], [0:27]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/slice[start_indices=(0, 0) limit_indices=(135168, 27) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=73}
  %slice.308 = f32[135168,108]{1,0} slice(f32[135168,150]{1,0} %param_1.522), slice={[0:135168], [30:138]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/slice[start_indices=(0, 30) limit_indices=(135168, 138) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=73}
  %bitcast.85 = f32[135168,27,4]{2,1,0} bitcast(f32[135168,108]{1,0} %slice.308), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reshape[new_sizes=(135168, 27, 4) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=37}
  %param_0.389 = f32[135168,27,4]{2,1,0} parameter(0)
  %cosine.1 = f32[135168,27,4]{2,1,0} cosine(f32[135168,27,4]{2,1,0} %param_0.389), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/cos" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=30}
  %multiply.354 = f32[135168,27,4]{2,1,0} multiply(f32[135168,27,4]{2,1,0} %bitcast.85, f32[135168,27,4]{2,1,0} %cosine.1), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=30}
  %slice.307 = f32[135168,27,2]{2,1,0} slice(f32[135168,27,4]{2,1,0} %multiply.354), slice={[0:135168], [0:27], [0:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/slice[start_indices=(0, 0, 0) limit_indices=(135168, 27, 2) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=31}
  %slice.306 = f32[135168,27,2]{2,1,0} slice(f32[135168,27,4]{2,1,0} %multiply.354), slice={[0:135168], [0:27], [2:4]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/slice[start_indices=(0, 0, 2) limit_indices=(135168, 27, 4) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=31}
  %add.411 = f32[135168,27,2]{2,1,0} add(f32[135168,27,2]{2,1,0} %slice.307, f32[135168,27,2]{2,1,0} %slice.306), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=32}
  %param_2.463 = f32[2]{0} parameter(2)
  %broadcast.1007 = f32[135168,27,2]{2,1,0} broadcast(f32[2]{0} %param_2.463), dimensions={2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %multiply.353 = f32[135168,27,2]{2,1,0} multiply(f32[135168,27,2]{2,1,0} %add.411, f32[135168,27,2]{2,1,0} %broadcast.1007), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %constant_585 = f32[] constant(0)
  %reduce.43 = f32[135168,27]{1,0} reduce(f32[135168,27,2]{2,1,0} %multiply.353, f32[] %constant_585), dimensions={2}, to_apply=%region_30.1722, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  ROOT %add.410 = f32[135168,27]{1,0} add(f32[135168,27]{1,0} %slice.309, f32[135168,27]{1,0} %reduce.43), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
}

%fused_computation.62 (param_0.141: f32[135168,128], param_1.523: f32[135168,128]) -> f32[135168,128] {
  %param_1.523 = f32[135168,128]{1,0} parameter(1)
  %constant_586 = f32[] constant(0)
  %broadcast.1009 = f32[135168,128]{1,0} broadcast(f32[] %constant_586), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gt" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  %compare.133 = pred[135168,128]{1,0} compare(f32[135168,128]{1,0} %param_1.523, f32[135168,128]{1,0} %broadcast.1009), direction=GT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gt" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  %param_0.141 = f32[135168,128]{1,0} parameter(0)
  ROOT %select.126 = f32[135168,128]{1,0} select(pred[135168,128]{1,0} %compare.133, f32[135168,128]{1,0} %param_0.141, f32[135168,128]{1,0} %broadcast.1009), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
}

%fused_computation.63 (param_0.143: f32[135168,128], param_1.524: f32[135168,128]) -> f32[135168,128] {
  %param_1.524 = f32[135168,128]{1,0} parameter(1)
  %constant_587 = f32[] constant(0)
  %broadcast.1011 = f32[135168,128]{1,0} broadcast(f32[] %constant_587), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gt" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  %compare.134 = pred[135168,128]{1,0} compare(f32[135168,128]{1,0} %param_1.524, f32[135168,128]{1,0} %broadcast.1011), direction=GT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gt" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=105}
  %param_0.143 = f32[135168,128]{1,0} parameter(0)
  ROOT %select.127 = f32[135168,128]{1,0} select(pred[135168,128]{1,0} %compare.134, f32[135168,128]{1,0} %param_0.143, f32[135168,128]{1,0} %broadcast.1011), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=105}
}

%fused_computation.64 (param_0.688: pred[4096,33,3], param_1.906: s32[4096,33,2], param_2.928: f32[4096,221], param_3.816: f32[135168,3], param_4.592: f32[4096,3], param_5.690: f32[4096], param_6.653: f32[4096,221]) -> f32[135168,3] {
  %param_0.688 = pred[4096,33,3]{1,0,2} parameter(0)
  %param_4.592 = f32[4096,3]{1,0} parameter(4)
  %constant_589 = f32[] constant(1)
  %broadcast.2346 = f32[4096,1]{1,0} broadcast(f32[] %constant_589), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_6.653 = f32[4096,221]{1,0} parameter(6)
  %slice.734 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_6.653), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(1,), start_index_map=(1,)) slice_sizes=(4096, 1) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=482}
  %subtract.406 = f32[4096,1]{1,0} subtract(f32[4096,1]{1,0} %broadcast.2346, f32[4096,1]{1,0} %slice.734), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %constant_1888 = f32[] constant(1e-08)
  %broadcast.2345 = f32[4096,1]{1,0} broadcast(f32[] %constant_1888), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.1072 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %subtract.406, f32[4096,1]{1,0} %broadcast.2345), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %bitcast.552 = f32[4096]{0} bitcast(f32[4096,1]{1,0} %add.1072), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_5.690 = f32[4096]{0} parameter(5)
  %broadcast.2344 = f32[4096]{0} broadcast(f32[] %constant_1888), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.1071 = f32[4096]{0} add(f32[4096]{0} %param_5.690, f32[4096]{0} %broadcast.2344), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  %divide.62 = f32[4096]{0} divide(f32[4096]{0} %bitcast.552, f32[4096]{0} %add.1071), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %constant_1715 = f32[] constant(0.000162760422)
  %broadcast.2175 = f32[4096]{0} broadcast(f32[] %constant_1715), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.577 = f32[4096]{0} multiply(f32[4096]{0} %divide.62, f32[4096]{0} %broadcast.2175), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %broadcast.2174 = f32[4096,3]{1,0} broadcast(f32[4096]{0} %multiply.577), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.576 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %param_4.592, f32[4096,3]{1,0} %broadcast.2174), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %broadcast.2173 = f32[4096,221,3]{1,0,2} broadcast(f32[4096,3]{1,0} %multiply.576), dimensions={0,2}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/broadcast_in_dim[shape=(4096, 221, 3) broadcast_dimensions=(0, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %param_2.928 = f32[4096,221]{1,0} parameter(2)
  %broadcast.1013 = f32[4096,221,3]{1,0,2} broadcast(f32[4096,221]{1,0} %param_2.928), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.357 = f32[4096,221,3]{1,0,2} multiply(f32[4096,221,3]{1,0,2} %broadcast.2173, f32[4096,221,3]{1,0,2} %broadcast.1013), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %param_1.906 = s32[4096,33,2]{1,0,2} parameter(1)
  %gather.48 = f32[4096,33,3]{1,0,2} gather(f32[4096,221,3]{1,0,2} %multiply.357, s32[4096,33,2]{1,0,2} %param_1.906), offset_dims={2}, collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=2, slice_sizes={1,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 1), start_index_map=(0, 1)) slice_sizes=(1, 1, 3) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %constant_588 = f32[] constant(0)
  %broadcast.1014 = f32[4096,33,3]{1,0,2} broadcast(f32[] %constant_588), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 33, 3) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %select.128 = f32[4096,33,3]{1,0,2} select(pred[4096,33,3]{1,0,2} %param_0.688, f32[4096,33,3]{1,0,2} %gather.48, f32[4096,33,3]{1,0,2} %broadcast.1014), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %bitcast.86 = f32[135168,3]{0,1} bitcast(f32[4096,33,3]{1,0,2} %select.128), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reshape[new_sizes=(135168, 3) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=440}
  %broadcast.1016 = f32[135168,3]{0,1} broadcast(f32[] %constant_589), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %param_3.816 = f32[135168,3]{0,1} parameter(3)
  %negate.43 = f32[135168,3]{0,1} negate(f32[135168,3]{0,1} %param_3.816), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/neg" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %exponential.24 = f32[135168,3]{0,1} exponential(f32[135168,3]{0,1} %negate.43), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/exp" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %add.921 = f32[135168,3]{0,1} add(f32[135168,3]{0,1} %exponential.24, f32[135168,3]{0,1} %broadcast.1016), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %divide.57 = f32[135168,3]{0,1} divide(f32[135168,3]{0,1} %broadcast.1016, f32[135168,3]{0,1} %add.921), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/div" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %subtract.66 = f32[135168,3]{0,1} subtract(f32[135168,3]{0,1} %broadcast.1016, f32[135168,3]{0,1} %divide.57), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/sub" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %multiply.356 = f32[135168,3]{0,1} multiply(f32[135168,3]{0,1} %bitcast.86, f32[135168,3]{0,1} %subtract.66), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(sigmoid)))/jit(jit_transpose(jvp(sigmoid)))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  ROOT %multiply.355 = f32[135168,3]{0,1} multiply(f32[135168,3]{0,1} %multiply.356, f32[135168,3]{0,1} %divide.57), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(sigmoid)))/jit(jit_transpose(jvp(sigmoid)))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
}

%fused_computation.66 (param_0.149: f32[4096,3], param_1.1011: f32[3,4096], param_2.1055: f32[4096,221], param_3.932: f32[4096]) -> f32[4096,3] {
  %param_1.1011 = f32[3,4096]{0,1} parameter(1)
  %bitcast.140 = f32[4096,3]{1,0} bitcast(f32[3,4096]{0,1} %param_1.1011), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %constant_1884 = f32[] constant(1)
  %broadcast.2340 = f32[4096,1]{1,0} broadcast(f32[] %constant_1884), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_2.1055 = f32[4096,221]{1,0} parameter(2)
  %slice.732 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_2.1055), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(1,), start_index_map=(1,)) slice_sizes=(4096, 1) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=482}
  %subtract.404 = f32[4096,1]{1,0} subtract(f32[4096,1]{1,0} %broadcast.2340, f32[4096,1]{1,0} %slice.732), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %constant_1883 = f32[] constant(1e-08)
  %broadcast.2339 = f32[4096,1]{1,0} broadcast(f32[] %constant_1883), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.1068 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %subtract.404, f32[4096,1]{1,0} %broadcast.2339), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %bitcast.550 = f32[4096]{0} bitcast(f32[4096,1]{1,0} %add.1068), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_3.932 = f32[4096]{0} parameter(3)
  %broadcast.2338 = f32[4096]{0} broadcast(f32[] %constant_1883), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.1067 = f32[4096]{0} add(f32[4096]{0} %param_3.932, f32[4096]{0} %broadcast.2338), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  %divide.60 = f32[4096]{0} divide(f32[4096]{0} %bitcast.550, f32[4096]{0} %add.1067), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %broadcast.898 = f32[4096,3]{1,0} broadcast(f32[4096]{0} %divide.60), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.360 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %bitcast.140, f32[4096,3]{1,0} %broadcast.898), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %bitcast.87 = f32[4096]{0} bitcast(f32[4096,1]{1,0} %slice.732), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=196}
  %broadcast.897 = f32[4096,3]{1,0} broadcast(f32[4096]{0} %bitcast.87), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=196}
  %add.412 = f32[4096,3]{1,0} add(f32[4096,3]{1,0} %multiply.360, f32[4096,3]{1,0} %broadcast.897), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=196}
  %param_0.149 = f32[4096,3]{1,0} parameter(0)
  ROOT %subtract.67 = f32[4096,3]{1,0} subtract(f32[4096,3]{1,0} %add.412, f32[4096,3]{1,0} %param_0.149), metadata={op_name="jit(training_step)/jit(main)/sub" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
}

%region_8.947 (Arg_0.948: f32[], Arg_1.949: f32[]) -> f32[] {
  %Arg_0.948 = f32[] parameter(0)
  %Arg_1.949 = f32[] parameter(1)
  ROOT %add.950 = f32[] add(f32[] %Arg_0.948, f32[] %Arg_1.949), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
}

%input_fused_computation_reduce.4 (param_0.390: f32[4096,221], param_1.527: s32[4096,33,2]) -> f32[4096] {
  %param_0.390 = f32[4096,221]{1,0} parameter(0)
  %param_1.527 = s32[4096,33,2]{1,0,2} parameter(1)
  %gather.49 = f32[4096,33]{1,0} gather(f32[4096,221]{1,0} %param_0.390, s32[4096,33,2]{1,0,2} %param_1.527), offset_dims={}, collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=2, slice_sizes={1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0, 1), start_index_map=(0, 1)) slice_sizes=(1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=472}
  %constant_592 = f32[] constant(0)
  ROOT %reduce.44 = f32[4096]{0} reduce(f32[4096,33]{1,0} %gather.49, f32[] %constant_592), dimensions={1}, to_apply=%region_8.947, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
}

%region_9.951 (Arg_0.952: f32[], Arg_1.953: f32[]) -> f32[] {
  %Arg_0.952 = f32[] parameter(0)
  %Arg_1.953 = f32[] parameter(1)
  ROOT %add.954 = f32[] add(f32[] %Arg_0.952, f32[] %Arg_1.953), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
}

%input_fused_computation_reduce.5 (param_0.393: f32[4096,221,3], param_1.530: f32[4096,221]) -> f32[3,4096] {
  %param_0.393 = f32[4096,221,3]{1,0,2} parameter(0)
  %param_1.530 = f32[4096,221]{1,0} parameter(1)
  %broadcast.1018 = f32[4096,221,3]{1,0,2} broadcast(f32[4096,221]{1,0} %param_1.530), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.361 = f32[4096,221,3]{1,0,2} multiply(f32[4096,221,3]{1,0,2} %param_0.393, f32[4096,221,3]{1,0,2} %broadcast.1018), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %bitcast.89 = f32[3,4096,221]{2,1,0} bitcast(f32[4096,221,3]{1,0,2} %multiply.361), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %constant_595 = f32[] constant(0)
  ROOT %reduce.45 = f32[3,4096]{0,1} reduce(f32[3,4096,221]{2,1,0} %bitcast.89, f32[] %constant_595), dimensions={2}, to_apply=%region_9.951
}

%region_7.943 (Arg_0.944: f32[], Arg_1.945: f32[]) -> f32[] {
  %Arg_0.944 = f32[] parameter(0)
  %Arg_1.945 = f32[] parameter(1)
  ROOT %add.946 = f32[] add(f32[] %Arg_0.944, f32[] %Arg_1.945), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.13 (param_0.159: s32[4096,33,2], param_1.240: pred[4096,33,3], param_2.907: f32[135168,3]) -> f32[4096,221,3] {
  %constant_596 = f32[] constant(0)
  %broadcast.902 = f32[4096,221,3]{1,0,2} broadcast(f32[] %constant_596), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 221, 3) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %param_0.159 = s32[4096,33,2]{1,0,2} parameter(0)
  %param_1.240 = pred[4096,33,3]{1,0,2} parameter(1)
  %constant_1677 = f32[] constant(1)
  %broadcast.2113 = f32[135168,3]{0,1} broadcast(f32[] %constant_1677), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %param_2.907 = f32[135168,3]{0,1} parameter(2)
  %negate.41 = f32[135168,3]{0,1} negate(f32[135168,3]{0,1} %param_2.907), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/neg" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %exponential.22 = f32[135168,3]{0,1} exponential(f32[135168,3]{0,1} %negate.41), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/exp" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %add.919 = f32[135168,3]{0,1} add(f32[135168,3]{0,1} %exponential.22, f32[135168,3]{0,1} %broadcast.2113), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %divide.55 = f32[135168,3]{0,1} divide(f32[135168,3]{0,1} %broadcast.2113, f32[135168,3]{0,1} %add.919), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/div" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %bitcast.90 = f32[4096,33,3]{1,0,2} bitcast(f32[135168,3]{0,1} %divide.55), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 3) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=440}
  %broadcast.1020 = f32[4096,33,3]{1,0,2} broadcast(f32[] %constant_596), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 33, 3) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %select.129 = f32[4096,33,3]{1,0,2} select(pred[4096,33,3]{1,0,2} %param_1.240, f32[4096,33,3]{1,0,2} %bitcast.90, f32[4096,33,3]{1,0,2} %broadcast.1020), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  ROOT %scatter.55 = f32[4096,221,3]{1,0,2} scatter(f32[4096,221,3]{1,0,2} %broadcast.902, s32[4096,33,2]{1,0,2} %param_0.159, f32[4096,33,3]{1,0,2} %select.129), update_window_dims={2}, inserted_window_dims={0,1}, scatter_dims_to_operand_dims={0,1}, index_vector_dim=2, to_apply=%region_7.943, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 1), scatter_dims_to_operand_dims=(0, 1)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%fused_computation.70 (param_0.689: u32[4096,221,3], param_1.911: s32[4096,33,2]) -> pred[4096,33,3] {
  %iota.199 = u32[4096,33,1]{2,1,0} iota(), iota_dimension=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/iota[dtype=uint32 shape=(135168,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %iota.198 = u32[4096,33,1]{2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/iota[dtype=uint32 shape=(135168,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %constant_1690 = u32[] constant(33), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %broadcast.2134 = u32[4096,33,1]{2,1,0} broadcast(u32[] %constant_1690), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %multiply.569 = u32[4096,33,1]{2,1,0} multiply(u32[4096,33,1]{2,1,0} %iota.198, u32[4096,33,1]{2,1,0} %broadcast.2134), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %add.929 = u32[4096,33,1]{2,1,0} add(u32[4096,33,1]{2,1,0} %iota.199, u32[4096,33,1]{2,1,0} %multiply.569), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %constant_1689 = u32[] constant(1)
  %broadcast.2132 = u32[4096,33,1]{2,1,0} broadcast(u32[] %constant_1689), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 33, 3) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %add.928 = u32[4096,33,1]{2,1,0} add(u32[4096,33,1]{2,1,0} %add.929, u32[4096,33,1]{2,1,0} %broadcast.2132), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %bitcast.540 = u32[4096,33]{1,0} bitcast(u32[4096,33,1]{2,1,0} %add.928), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %broadcast.2130 = u32[4096,33,3]{1,0,2} broadcast(u32[4096,33]{1,0} %bitcast.540), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %param_0.689 = u32[4096,221,3]{1,0,2} parameter(0)
  %param_1.911 = s32[4096,33,2]{1,0,2} parameter(1)
  %gather.50 = u32[4096,33,3]{1,0,2} gather(u32[4096,221,3]{1,0,2} %param_0.689, s32[4096,33,2]{1,0,2} %param_1.911), offset_dims={2}, collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=2, slice_sizes={1,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 1), start_index_map=(0, 1)) slice_sizes=(1, 1, 3) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  ROOT %compare.135 = pred[4096,33,3]{1,0,2} compare(u32[4096,33,3]{1,0,2} %broadcast.2130, u32[4096,33,3]{1,0,2} %gather.50), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%region_6.940 (Arg_0.941: u32[], Arg_1.942: u32[]) -> u32[] {
  %Arg_0.941 = u32[] parameter(0)
  ROOT %Arg_1.942 = u32[] parameter(1)
}

%input_fused_computation_scatter.14 (param_0.163: s32[4096,33,2]) -> u32[4096,221,3] {
  %constant_457 = u32[] constant(0)
  %broadcast.903 = u32[4096,221,3]{1,0,2} broadcast(u32[] %constant_457), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 221, 3) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %param_0.163 = s32[4096,33,2]{1,0,2} parameter(0)
  %iota.195 = u32[4096,33,1]{2,1,0} iota(), iota_dimension=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/iota[dtype=uint32 shape=(135168,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %iota.194 = u32[4096,33,1]{2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/iota[dtype=uint32 shape=(135168,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %constant_1685 = u32[] constant(33), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %broadcast.2122 = u32[4096,33,1]{2,1,0} broadcast(u32[] %constant_1685), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %multiply.567 = u32[4096,33,1]{2,1,0} multiply(u32[4096,33,1]{2,1,0} %iota.194, u32[4096,33,1]{2,1,0} %broadcast.2122), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %add.925 = u32[4096,33,1]{2,1,0} add(u32[4096,33,1]{2,1,0} %iota.195, u32[4096,33,1]{2,1,0} %multiply.567), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %constant_1684 = u32[] constant(1)
  %broadcast.2120 = u32[4096,33,1]{2,1,0} broadcast(u32[] %constant_1684), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 33, 3) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %add.924 = u32[4096,33,1]{2,1,0} add(u32[4096,33,1]{2,1,0} %add.925, u32[4096,33,1]{2,1,0} %broadcast.2120), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %bitcast.538 = u32[4096,33]{1,0} bitcast(u32[4096,33,1]{2,1,0} %add.924), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %broadcast.2119 = u32[4096,33,3]{1,0,2} broadcast(u32[4096,33]{1,0} %bitcast.538), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  ROOT %scatter.56 = u32[4096,221,3]{1,0,2} scatter(u32[4096,221,3]{1,0,2} %broadcast.903, s32[4096,33,2]{1,0,2} %param_0.163, u32[4096,33,3]{1,0,2} %broadcast.2119), update_window_dims={2}, inserted_window_dims={0,1}, scatter_dims_to_operand_dims={0,1}, index_vector_dim=2, to_apply=%region_6.940, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/scatter[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 1), scatter_dims_to_operand_dims=(0, 1)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%fused_computation.72 (param_0.168: f32[135168,27], param_1.536: f32[2], param_2.472: f32[135168,27,4], param_3.366: f32[4096,3]) -> f32[135168,150] {
  %param_0.168 = f32[135168,27]{1,0} parameter(0)
  %param_3.366 = f32[4096,3]{1,0} parameter(3)
  %broadcast.910 = f32[1,4096,33,1,1,3]{5,4,3,2,1,0} broadcast(f32[4096,3]{1,0} %param_3.366), dimensions={1,5}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(1, 4096, 33, 1, 1, 3) broadcast_dimensions=(0, 1, 3, 4, 5)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=435}
  %bitcast.94 = f32[135168,3]{1,0} bitcast(f32[1,4096,33,1,1,3]{5,4,3,2,1,0} %broadcast.910), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(135168, 3) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=435}
  %param_2.472 = f32[135168,27,4]{2,1,0} parameter(2)
  %sine.3 = f32[135168,27,4]{2,1,0} sine(f32[135168,27,4]{2,1,0} %param_2.472), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sin" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=30}
  %bitcast.93 = f32[135168,108]{1,0} bitcast(f32[135168,27,4]{2,1,0} %sine.3), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(135168, 108) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=37}
  %broadcast.909 = f32[135168,3,2]{2,1,0} broadcast(f32[135168,3]{1,0} %bitcast.94), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %param_1.536 = f32[2]{0} parameter(1)
  %broadcast.908 = f32[135168,3,2]{2,1,0} broadcast(f32[2]{0} %param_1.536), dimensions={2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %multiply.363 = f32[135168,3,2]{2,1,0} multiply(f32[135168,3,2]{2,1,0} %broadcast.909, f32[135168,3,2]{2,1,0} %broadcast.908), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %constant_600 = f32[] constant(1.57079637)
  %broadcast.907 = f32[135168,3,2]{2,1,0} broadcast(f32[] %constant_600), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=32}
  %add.420 = f32[135168,3,2]{2,1,0} add(f32[135168,3,2]{2,1,0} %multiply.363, f32[135168,3,2]{2,1,0} %broadcast.907), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=32}
  %concatenate.109 = f32[135168,3,4]{2,1,0} concatenate(f32[135168,3,2]{2,1,0} %multiply.363, f32[135168,3,2]{2,1,0} %add.420), dimensions={2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=2]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=31}
  %sine.2 = f32[135168,3,4]{2,1,0} sine(f32[135168,3,4]{2,1,0} %concatenate.109), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sin" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=30}
  %bitcast.92 = f32[135168,12]{1,0} bitcast(f32[135168,3,4]{2,1,0} %sine.2), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(135168, 12) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=37}
  ROOT %concatenate.108 = f32[135168,150]{1,0} concatenate(f32[135168,27]{1,0} %param_0.168, f32[135168,3]{1,0} %bitcast.94, f32[135168,108]{1,0} %bitcast.93, f32[135168,12]{1,0} %bitcast.92), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=73}
}

%fused_computation.73 (param_0.398: f32[135168,27], param_1.538: f32[2]) -> f32[135168,27,4] {
  %param_0.398 = f32[135168,27]{1,0} parameter(0)
  %broadcast.912 = f32[135168,27,2]{2,1,0} broadcast(f32[135168,27]{1,0} %param_0.398), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %param_1.538 = f32[2]{0} parameter(1)
  %broadcast.1028 = f32[135168,27,2]{2,1,0} broadcast(f32[2]{0} %param_1.538), dimensions={2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %multiply.364 = f32[135168,27,2]{2,1,0} multiply(f32[135168,27,2]{2,1,0} %broadcast.912, f32[135168,27,2]{2,1,0} %broadcast.1028), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %constant_601 = f32[] constant(1.57079637)
  %broadcast.911 = f32[135168,27,2]{2,1,0} broadcast(f32[] %constant_601), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=32}
  %add.421 = f32[135168,27,2]{2,1,0} add(f32[135168,27,2]{2,1,0} %multiply.364, f32[135168,27,2]{2,1,0} %broadcast.911), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=32}
  ROOT %concatenate.110 = f32[135168,27,4]{2,1,0} concatenate(f32[135168,27,2]{2,1,0} %multiply.364, f32[135168,27,2]{2,1,0} %add.421), dimensions={2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=2]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=31}
}

%fused_computation.74 () -> f32[2] {
  %constant_602 = f32[] constant(2)
  %broadcast.913 = f32[2]{0} broadcast(f32[] %constant_602), dimensions={}
  %iota.85 = f32[2]{0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/iota[dtype=float32 shape=(2,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=26}
  ROOT %power.5 = f32[2]{0} power(f32[2]{0} %broadcast.913, f32[2]{0} %iota.85), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/pow" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=26}
}

%fused_computation.75 (param_0.685: f32[3,48,128], param_1.905: f32[3,3,4096,33], param_2.905: f32[3,48,128,128]) -> f32[135168,144] {
  %constant_1329 = f32[] constant(1)
  %broadcast.1745 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1329), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.905 = f32[3,3,4096,33]{3,2,0,1} parameter(1)
  %slice.450 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.905), slice={[0:3], [0:1], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1, 2, 3), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(3, 1, 4096, 33) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=165}
  %bitcast.296 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.450), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.235 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.296), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.260 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.296, f32[3,4096,33]{2,1,0} %floor.235), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.259 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1745, f32[3,4096,33]{2,1,0} %subtract.260), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1744 = f32[3,48,4096,33]{2,1,0,3} broadcast(f32[3,4096,33]{2,1,0} %subtract.259), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.685 = f32[3,48,128]{2,1,0} parameter(0)
  %iota.159 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1325 = s32[] constant(0)
  %broadcast.1743 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1325), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.138 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.235), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1323 = s32[] constant(127)
  %broadcast.1742 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1323), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.127 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1743, s32[3,4096,33]{2,1,0} %convert.138, s32[3,4096,33]{2,1,0} %broadcast.1742), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.253 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.127, s32[3,4096,33]{2,1,0} %broadcast.1743), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1322 = s32[] constant(128)
  %broadcast.1741 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1322), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.769 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.127, s32[3,4096,33]{2,1,0} %broadcast.1741), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.300 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.253, s32[3,4096,33]{2,1,0} %add.769, s32[3,4096,33]{2,1,0} %clamp.127), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.294 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.300), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.210 = s32[3,4096,33,2]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.159, s32[3,4096,33,1]{2,1,0,3} %bitcast.294), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.103 = f32[3,48,4096,33]{2,1,0,3} gather(f32[3,48,128]{2,1,0} %param_0.685, s32[3,4096,33,2]{2,1,0,3} %concatenate.210), offset_dims={1}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,48,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 48, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.516 = f32[3,48,4096,33]{2,1,0,3} multiply(f32[3,48,4096,33]{2,1,0,3} %broadcast.1744, f32[3,48,4096,33]{2,1,0,3} %gather.103), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.33 = f32[3,48,4096,33]{3,2,1,0} copy(f32[3,48,4096,33]{2,1,0,3} %multiply.516), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1740 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %subtract.260), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.158 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1324 = s32[] constant(1)
  %broadcast.1738 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1324), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.768 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.138, s32[3,4096,33]{2,1,0} %broadcast.1738), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.126 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1743, s32[3,4096,33]{2,1,0} %add.768, s32[3,4096,33]{2,1,0} %broadcast.1742), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.252 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.126, s32[3,4096,33]{2,1,0} %broadcast.1743), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.767 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.126, s32[3,4096,33]{2,1,0} %broadcast.1741), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.299 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.252, s32[3,4096,33]{2,1,0} %add.767, s32[3,4096,33]{2,1,0} %clamp.126), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.291 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.299), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.209 = s32[3,4096,33,2]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.158, s32[3,4096,33,1]{3,2,1,0} %bitcast.291), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.102 = f32[3,48,4096,33]{3,2,1,0} gather(f32[3,48,128]{2,1,0} %param_0.685, s32[3,4096,33,2]{3,2,1,0} %concatenate.209), offset_dims={1}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,48,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 48, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.515 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.1740, f32[3,48,4096,33]{3,2,1,0} %gather.102), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.766 = f32[3,48,4096,33]{3,2,1,0} add(f32[3,48,4096,33]{3,2,1,0} %copy.33, f32[3,48,4096,33]{3,2,1,0} %multiply.515), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.650 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.905), slice={[0:3], [1:2], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 1, 4096, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.536 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.650), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.382 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.536), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.358 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.536, f32[3,4096,33]{2,1,0} %floor.382), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.357 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1745, f32[3,4096,33]{2,1,0} %subtract.358), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.649 = f32[3,1,4096,33]{3,2,0,1} slice(f32[3,3,4096,33]{3,2,0,1} %param_1.905), slice={[0:3], [2:3], [0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 1, 0, 0) limit_indices=(3, 2, 4096, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.535 = f32[3,4096,33]{2,1,0} bitcast(f32[3,1,4096,33]{3,2,0,1} %slice.649), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.381 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.535), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.356 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.535, f32[3,4096,33]{2,1,0} %floor.381), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.355 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1745, f32[3,4096,33]{2,1,0} %subtract.356), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.565 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.357, f32[3,4096,33]{2,1,0} %subtract.355), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2110 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.565), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.905 = f32[3,48,128,128]{3,2,1,0} parameter(2)
  %iota.191 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.218 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.382), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.207 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1743, s32[3,4096,33]{2,1,0} %convert.218, s32[3,4096,33]{2,1,0} %broadcast.1742), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.335 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.207, s32[3,4096,33]{2,1,0} %broadcast.1743), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.917 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.207, s32[3,4096,33]{2,1,0} %broadcast.1741), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.383 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.335, s32[3,4096,33]{2,1,0} %add.917, s32[3,4096,33]{2,1,0} %clamp.207), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.533 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.383), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.217 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.381), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.206 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1743, s32[3,4096,33]{2,1,0} %convert.217, s32[3,4096,33]{2,1,0} %broadcast.1742), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.334 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.206, s32[3,4096,33]{2,1,0} %broadcast.1743), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.916 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.206, s32[3,4096,33]{2,1,0} %broadcast.1741), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.381 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.334, s32[3,4096,33]{2,1,0} %add.916, s32[3,4096,33]{2,1,0} %clamp.206), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.531 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.381), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.243 = s32[3,4096,33,3]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.191, s32[3,4096,33,1]{2,1,0,3} %bitcast.533, s32[3,4096,33,1]{2,1,0,3} %bitcast.531), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.119 = f32[3,48,4096,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.905, s32[3,4096,33,3]{2,1,0,3} %concatenate.243), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.564 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.2110, f32[3,48,4096,33]{3,2,1,0} %gather.119), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.563 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.357, f32[3,4096,33]{2,1,0} %subtract.356), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2102 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.563), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.190 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.527 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.383), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.914 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.217, s32[3,4096,33]{2,1,0} %broadcast.1738), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.204 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1743, s32[3,4096,33]{2,1,0} %add.914, s32[3,4096,33]{2,1,0} %broadcast.1742), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.332 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.204, s32[3,4096,33]{2,1,0} %broadcast.1743), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.913 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.204, s32[3,4096,33]{2,1,0} %broadcast.1741), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.378 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.332, s32[3,4096,33]{2,1,0} %add.913, s32[3,4096,33]{2,1,0} %clamp.204), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.525 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.378), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.242 = s32[3,4096,33,3]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.190, s32[3,4096,33,1]{3,2,1,0} %bitcast.527, s32[3,4096,33,1]{3,2,1,0} %bitcast.525), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.118 = f32[3,48,4096,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.905, s32[3,4096,33,3]{3,2,1,0} %concatenate.242), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.562 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.2102, f32[3,48,4096,33]{3,2,1,0} %gather.118), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.912 = f32[3,48,4096,33]{3,2,1,0} add(f32[3,48,4096,33]{3,2,1,0} %multiply.564, f32[3,48,4096,33]{3,2,1,0} %multiply.562), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.561 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.358, f32[3,4096,33]{2,1,0} %subtract.355), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2093 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.561), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.189 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.911 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.218, s32[3,4096,33]{2,1,0} %broadcast.1738), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.203 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1743, s32[3,4096,33]{2,1,0} %add.911, s32[3,4096,33]{2,1,0} %broadcast.1742), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.331 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.203, s32[3,4096,33]{2,1,0} %broadcast.1743), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.910 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.203, s32[3,4096,33]{2,1,0} %broadcast.1741), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.377 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.331, s32[3,4096,33]{2,1,0} %add.910, s32[3,4096,33]{2,1,0} %clamp.203), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.521 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.377), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.241 = s32[3,4096,33,3]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.189, s32[3,4096,33,1]{2,1,0,3} %bitcast.521, s32[3,4096,33,1]{2,1,0,3} %bitcast.531), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.117 = f32[3,48,4096,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.905, s32[3,4096,33,3]{2,1,0,3} %concatenate.241), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.560 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.2093, f32[3,48,4096,33]{3,2,1,0} %gather.117), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.908 = f32[3,48,4096,33]{3,2,1,0} add(f32[3,48,4096,33]{3,2,1,0} %add.912, f32[3,48,4096,33]{3,2,1,0} %multiply.560), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.559 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.358, f32[3,4096,33]{2,1,0} %subtract.356), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2085 = f32[3,48,4096,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.559), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.188 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.515 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.377), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.240 = s32[3,4096,33,3]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.188, s32[3,4096,33,1]{3,2,1,0} %bitcast.515, s32[3,4096,33,1]{3,2,1,0} %bitcast.525), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.116 = f32[3,48,4096,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.905, s32[3,4096,33,3]{3,2,1,0} %concatenate.240), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.558 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %broadcast.2085, f32[3,48,4096,33]{3,2,1,0} %gather.116), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.903 = f32[3,48,4096,33]{3,2,1,0} add(f32[3,48,4096,33]{3,2,1,0} %add.908, f32[3,48,4096,33]{3,2,1,0} %multiply.558), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.365 = f32[3,48,4096,33]{3,2,1,0} multiply(f32[3,48,4096,33]{3,2,1,0} %add.766, f32[3,48,4096,33]{3,2,1,0} %add.903), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %bitcast.96 = f32[144,4096,33]{2,1,0} bitcast(f32[3,48,4096,33]{3,2,1,0} %multiply.365), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/reshape[new_sizes=(144, 4096, 33) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %transpose.11 = f32[4096,33,144]{1,0,2} transpose(f32[144,4096,33]{2,1,0} %bitcast.96), dimensions={1,2,0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/transpose[permutation=(1, 2, 0)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
  ROOT %bitcast.95 = f32[135168,144]{0,1} bitcast(f32[4096,33,144]{1,0,2} %transpose.11), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(135168, 144) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
}

%fused_computation.96 (param_0.425: s32[3,1], param_1.581: s32[3,1], param_2.523: f32[3,4096,221], param_3.408: s32[4096,33,2]) -> f32[3,3,4096,33] {
  %param_2.523 = f32[3,4096,221]{2,1,0} parameter(2)
  %param_3.408 = s32[4096,33,2]{1,0,2} parameter(3)
  %gather.59 = f32[3,4096,33]{2,1,0} gather(f32[3,4096,221]{2,1,0} %param_2.523, s32[4096,33,2]{1,0,2} %param_3.408), offset_dims={0}, collapsed_slice_dims={1,2}, start_index_map={1,2}, index_vector_dim=2, slice_sizes={3,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(1, 2), start_index_map=(1, 2)) slice_sizes=(3, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %bitcast.110 = f32[1,3,4096,33]{3,2,0,1} bitcast(f32[3,4096,33]{2,1,0} %gather.59), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 33) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %param_1.581 = s32[3,1]{1,0} parameter(1)
  %gather.58 = f32[3,4096,33]{2,1,0} gather(f32[3,4096,33]{2,1,0} %gather.59, s32[3,1]{1,0} %param_1.581), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,4096,33}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 4096, 33) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=58}
  %bitcast.109 = f32[1,3,4096,33]{3,2,0,1} bitcast(f32[3,4096,33]{2,1,0} %gather.58), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 33) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %param_0.425 = s32[3,1]{1,0} parameter(0)
  %gather.57 = f32[3,4096,33]{2,1,0} gather(f32[3,4096,33]{2,1,0} %gather.59, s32[3,1]{1,0} %param_0.425), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,4096,33}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 4096, 33) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=59}
  %bitcast.108 = f32[1,3,4096,33]{3,2,0,1} bitcast(f32[3,4096,33]{2,1,0} %gather.57), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 33) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %concatenate.117 = f32[3,3,4096,33]{3,2,0,1} concatenate(f32[1,3,4096,33]{3,2,0,1} %bitcast.110, f32[1,3,4096,33]{3,2,0,1} %bitcast.109, f32[1,3,4096,33]{3,2,0,1} %bitcast.108), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %constant_629 = f32[] constant(1)
  %broadcast.920 = f32[3,3,4096,33]{3,2,0,1} broadcast(f32[] %constant_629), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=159}
  %add.437 = f32[3,3,4096,33]{3,2,0,1} add(f32[3,3,4096,33]{3,2,0,1} %concatenate.117, f32[3,3,4096,33]{3,2,0,1} %broadcast.920), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=159}
  %constant_628 = f32[] constant(63.5)
  %broadcast.919 = f32[3,3,4096,33]{3,2,0,1} broadcast(f32[] %constant_628), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=162}
  ROOT %multiply.376 = f32[3,3,4096,33]{3,2,0,1} multiply(f32[3,3,4096,33]{3,2,0,1} %add.437, f32[3,3,4096,33]{3,2,0,1} %broadcast.919), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=162}
}

%fused_computation.97 (param_0.426: s32[4096,221]) -> s32[4096,33,2] {
  %iota.92 = s32[4096,33,1]{1,0,2} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 33) broadcast_dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %param_0.426 = s32[4096,221]{1,0} parameter(0)
  %slice.313 = s32[4096,33]{1,0} slice(s32[4096,221]{1,0} %param_0.426), slice={[0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(4096, 33) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_634 = s32[] constant(0)
  %broadcast.922 = s32[4096,33]{1,0} broadcast(s32[] %constant_634), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %compare.142 = pred[4096,33]{1,0} compare(s32[4096,33]{1,0} %slice.313, s32[4096,33]{1,0} %broadcast.922), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %constant_462 = s32[] constant(221)
  %broadcast.921 = s32[4096,33]{1,0} broadcast(s32[] %constant_462), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %add.438 = s32[4096,33]{1,0} add(s32[4096,33]{1,0} %slice.313, s32[4096,33]{1,0} %broadcast.921), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %select.136 = s32[4096,33]{1,0} select(pred[4096,33]{1,0} %compare.142, s32[4096,33]{1,0} %add.438, s32[4096,33]{1,0} %slice.313), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %bitcast.111 = s32[4096,33,1]{1,0,2} bitcast(s32[4096,33]{1,0} %select.136), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 33, 1) broadcast_dimensions=(0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  ROOT %concatenate.118 = s32[4096,33,2]{1,0,2} concatenate(s32[4096,33,1]{1,0,2} %iota.92, s32[4096,33,1]{1,0,2} %bitcast.111), dimensions={2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=2]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
}

%fused_computation.98 (param_0.747: f32[221], param_1.998: pred[4096], param_2.1038: f32[4096], param_3.912: f32[4096], param_4.687: f32[4096,221], param_5.657: f32[4096,110], param_6.623: f32[4096,110]) -> (f32[4096,221], f32[4096,221], f32[4096,221]) {
  %param_0.747 = f32[221]{0} parameter(0)
  %broadcast.924 = f32[4096,221]{1,0} broadcast(f32[221]{0} %param_0.747), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_1247_clone_1 = f32[] constant(1)
  %broadcast.1629.clone.1 = f32[4096,221]{1,0} broadcast(f32[] %constant_1247_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %param_4.687 = f32[4096,221]{1,0} parameter(4)
  %negate.33.clone.1 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_4.687), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_1.998 = pred[4096]{0} parameter(1)
  %param_3.912 = f32[4096]{0} parameter(3)
  %param_2.1038 = f32[4096]{0} parameter(2)
  %constant_1204_clone_1 = f32[] constant(0.001)
  %broadcast.1628.clone.1 = f32[4096]{0} broadcast(f32[] %constant_1204_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.695.clone.1 = f32[4096]{0} add(f32[4096]{0} %param_2.1038, f32[4096]{0} %broadcast.1628.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.83.clone.1 = f32[4096]{0} maximum(f32[4096]{0} %param_3.912, f32[4096]{0} %add.695.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.271.clone.1 = f32[4096]{0} select(pred[4096]{0} %param_1.998, f32[4096]{0} %maximum.83.clone.1, f32[4096]{0} %broadcast.1628.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %constant_1203_clone_1 = f32[] constant(0)
  %broadcast.1627.clone.1 = f32[4096]{0} broadcast(f32[] %constant_1203_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.270.clone.1 = f32[4096]{0} select(pred[4096]{0} %param_1.998, f32[4096]{0} %param_2.1038, f32[4096]{0} %broadcast.1627.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.232.clone.1 = f32[4096]{0} subtract(f32[4096]{0} %select.271.clone.1, f32[4096]{0} %select.270.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1202_clone_1 = f32[] constant(0.00452488707)
  %broadcast.1626.clone.1 = f32[4096]{0} broadcast(f32[] %constant_1202_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.495.clone.1 = f32[4096]{0} multiply(f32[4096]{0} %subtract.232.clone.1, f32[4096]{0} %broadcast.1626.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.1625.clone.1 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.495.clone.1), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.494.clone.1 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.33.clone.1, f32[4096,221]{1,0} %broadcast.1625.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %exponential.18.clone.1 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %multiply.494.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %subtract.231.clone.1 = f32[4096,221]{1,0} subtract(f32[4096,221]{1,0} %broadcast.1629.clone.1, f32[4096,221]{1,0} %exponential.18.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %slice.316.clone.1 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %multiply.494.clone.1), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_5.657 = f32[4096,110]{1,0} parameter(5)
  %param_6.623 = f32[4096,110]{1,0} parameter(6)
  %add.441.clone.1 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %param_5.657, f32[4096,110]{1,0} %param_6.623), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.315.clone.1 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %multiply.494.clone.1), slice={[0:4096:1], [2:221:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.440.clone.1 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %add.441.clone.1, f32[4096,110]{1,0} %slice.315.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.119.clone.1 = f32[4096,111]{1,0} concatenate(f32[4096,1]{1,0} %slice.316.clone.1, f32[4096,110]{1,0} %add.440.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.115.clone.1 = f32[4096,221]{1,0} pad(f32[4096,111]{1,0} %concatenate.119.clone.1, f32[] %constant_1203_clone_1), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.114.clone.1 = f32[4096,221]{1,0} pad(f32[4096,110]{1,0} %add.441.clone.1, f32[] %constant_1203_clone_1), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.439.clone.1 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %pad.115.clone.1, f32[4096,221]{1,0} %pad.114.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %exponential.13.clone.1 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %add.439.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.410.clone.1 = f32[4096,220]{1,0} slice(f32[4096,221]{1,0} %exponential.13.clone.1), slice={[0:4096], [0:220]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(4096, 220) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=286}
  %pad.144.clone.1 = f32[4096,221]{1,0} pad(f32[4096,220]{1,0} %slice.410.clone.1, f32[] %constant_1247_clone_1), padding=0_0x1_0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=283}
  %multiply.491.clone.1 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %subtract.231.clone.1, f32[4096,221]{1,0} %pad.144.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=280}
  %log.4 = f32[4096,221]{1,0} log(f32[4096,221]{1,0} %multiply.491.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/log" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %subtract.70 = f32[4096,221]{1,0} subtract(f32[4096,221]{1,0} %broadcast.924, f32[4096,221]{1,0} %log.4), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  ROOT %tuple.76 = (f32[4096,221]{1,0}, f32[4096,221]{1,0}, f32[4096,221]{1,0}) tuple(f32[4096,221]{1,0} %subtract.70, f32[4096,221]{1,0} %multiply.491.clone.1, f32[4096,221]{1,0} %exponential.13.clone.1)
}

%fused_computation.105 (param_0.573: f32[4096,3], param_1.798: f32[4096,13]) -> f32[4096,6] {
  %param_1.798 = f32[4096,13]{1,0} parameter(1)
  %slice.388 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_1.798), slice={[0:4096:1], [0:12:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 12) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.387 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_1.798), slice={[0:4096:1], [1:13:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 13) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.714 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %slice.388, f32[4096,6]{1,0} %slice.387), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.328 = f32[4096,1]{1,0} slice(f32[4096,6]{1,0} %add.714), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_0.573 = f32[4096,3]{1,0} parameter(0)
  %slice.327 = f32[4096,2]{1,0} slice(f32[4096,3]{1,0} %param_0.573), slice={[0:4096], [0:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 2) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.326 = f32[4096,2]{1,0} slice(f32[4096,6]{1,0} %add.714), slice={[0:4096:1], [2:6:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 6) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.453 = f32[4096,2]{1,0} add(f32[4096,2]{1,0} %slice.327, f32[4096,2]{1,0} %slice.326), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.124 = f32[4096,3]{1,0} concatenate(f32[4096,1]{1,0} %slice.328, f32[4096,2]{1,0} %add.453), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_644 = f32[] constant(0)
  ROOT %pad.122 = f32[4096,6]{1,0} pad(f32[4096,3]{1,0} %concatenate.124, f32[] %constant_644), padding=0_0_0x0_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.106 (param_0.576: f32[4096,13]) -> f32[4096,3] {
  %param_0.576 = f32[4096,13]{1,0} parameter(0)
  %slice.392 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_0.576), slice={[0:4096:1], [0:12:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 12) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.391 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_0.576), slice={[0:4096:1], [1:13:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 13) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.716 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %slice.392, f32[4096,6]{1,0} %slice.391), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.334 = f32[4096,3]{1,0} slice(f32[4096,6]{1,0} %add.716), slice={[0:4096:1], [0:5:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 5) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.333 = f32[4096,3]{1,0} slice(f32[4096,6]{1,0} %add.716), slice={[0:4096:1], [1:6:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 6) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.457 = f32[4096,3]{1,0} add(f32[4096,3]{1,0} %slice.334, f32[4096,3]{1,0} %slice.333), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.332 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.457), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.331 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.457), slice={[0:4096:1], [0:2:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 2) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.330 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.457), slice={[0:4096:1], [1:3:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 3) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.456 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %slice.331, f32[4096,1]{1,0} %slice.330), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.329 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.457), slice={[0:4096:1], [2:3:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 3) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.455 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %add.456, f32[4096,1]{1,0} %slice.329), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.125 = f32[4096,2]{1,0} concatenate(f32[4096,1]{1,0} %slice.332, f32[4096,1]{1,0} %add.455), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_645 = f32[] constant(0)
  %pad.124 = f32[4096,3]{1,0} pad(f32[4096,2]{1,0} %concatenate.125, f32[] %constant_645), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.123 = f32[4096,3]{1,0} pad(f32[4096,1]{1,0} %add.456, f32[] %constant_645), padding=0_0x1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %add.454 = f32[4096,3]{1,0} add(f32[4096,3]{1,0} %pad.124, f32[4096,3]{1,0} %pad.123), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.108 (param_0.571: f32[4096,55]) -> f32[4096,13] {
  %param_0.571 = f32[4096,55]{1,0} parameter(0)
  %slice.384 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_0.571), slice={[0:4096:1], [0:54:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 54) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.383 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_0.571), slice={[0:4096:1], [1:55:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 55) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.711 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %slice.384, f32[4096,27]{1,0} %slice.383), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.338 = f32[4096,13]{1,0} slice(f32[4096,27]{1,0} %add.711), slice={[0:4096:1], [0:26:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 26) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.337 = f32[4096,13]{1,0} slice(f32[4096,27]{1,0} %add.711), slice={[0:4096:1], [1:27:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 27) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %add.459 = f32[4096,13]{1,0} add(f32[4096,13]{1,0} %slice.338, f32[4096,13]{1,0} %slice.337), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.110 (param_0.567: pred[4096], param_1.792: f32[4096], param_2.758: f32[4096], param_3.612: f32[4096,221]) -> f32[4096,55] {
  %param_3.612 = f32[4096,221]{1,0} parameter(3)
  %negate.39 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_3.612), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_0.567 = pred[4096]{0} parameter(0)
  %param_2.758 = f32[4096]{0} parameter(2)
  %param_1.792 = f32[4096]{0} parameter(1)
  %constant_1228 = f32[] constant(0.001)
  %broadcast.1664 = f32[4096]{0} broadcast(f32[] %constant_1228), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.707 = f32[4096]{0} add(f32[4096]{0} %param_1.792, f32[4096]{0} %broadcast.1664), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.89 = f32[4096]{0} maximum(f32[4096]{0} %param_2.758, f32[4096]{0} %add.707), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.283 = f32[4096]{0} select(pred[4096]{0} %param_0.567, f32[4096]{0} %maximum.89, f32[4096]{0} %broadcast.1664), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %constant_1227 = f32[] constant(0)
  %broadcast.1663 = f32[4096]{0} broadcast(f32[] %constant_1227), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.282 = f32[4096]{0} select(pred[4096]{0} %param_0.567, f32[4096]{0} %param_1.792, f32[4096]{0} %broadcast.1663), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.240 = f32[4096]{0} subtract(f32[4096]{0} %select.283, f32[4096]{0} %select.282), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1226 = f32[] constant(0.00452488707)
  %broadcast.1662 = f32[4096]{0} broadcast(f32[] %constant_1226), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.508 = f32[4096]{0} multiply(f32[4096]{0} %subtract.240, f32[4096]{0} %broadcast.1662), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.1661 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.508), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.507 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.39, f32[4096,221]{1,0} %broadcast.1661), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %slice.376 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %multiply.507), slice={[0:4096:1], [0:220:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 220) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.375 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %multiply.507), slice={[0:4096:1], [1:221:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.706 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %slice.376, f32[4096,110]{1,0} %slice.375), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.342 = f32[4096,55]{1,0} slice(f32[4096,110]{1,0} %add.706), slice={[0:4096:1], [0:109:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 109) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.341 = f32[4096,55]{1,0} slice(f32[4096,110]{1,0} %add.706), slice={[0:4096:1], [1:110:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 110) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %add.461 = f32[4096,55]{1,0} add(f32[4096,55]{1,0} %slice.342, f32[4096,55]{1,0} %slice.341), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.114 (param_0.547: f32[4096,221]) -> f32[4096,221] {
  %param_0.547 = f32[4096,221]{1,0} parameter(0)
  %constant_1168 = f32[] constant(10)
  %broadcast.1587 = f32[4096,221]{1,0} broadcast(f32[] %constant_1168), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %add.685 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %param_0.547, f32[4096,221]{1,0} %broadcast.1587), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.143 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %add.685, f32[4096,221]{1,0} %add.685), direction=NE, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/ne" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %constant_647 = f32[] constant(0)
  %broadcast.1071 = f32[4096,221]{1,0} broadcast(f32[] %constant_647), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %maximum.61 = f32[4096,221]{1,0} maximum(f32[4096,221]{1,0} %add.685, f32[4096,221]{1,0} %broadcast.1071), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %abs.3 = f32[4096,221]{1,0} abs(f32[4096,221]{1,0} %add.685), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/abs" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %negate.22 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %abs.3), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %exponential.14 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %negate.22), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %log-plus-one.3 = f32[4096,221]{1,0} log-plus-one(f32[4096,221]{1,0} %exponential.14), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/log1p" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %add.463 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %maximum.61, f32[4096,221]{1,0} %log-plus-one.3), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  ROOT %select.137 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.143, f32[4096,221]{1,0} %add.685, f32[4096,221]{1,0} %add.463), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
}

%region_4.507 (Arg_0.508: f32[], Arg_1.509: f32[]) -> f32[] {
  %Arg_0.508 = f32[] parameter(0)
  %Arg_1.509 = f32[] parameter(1)
  ROOT %add.510 = f32[] add(f32[] %Arg_0.508, f32[] %Arg_1.509), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
}

%input_fused_computation_reduce.6 (param_0.438: f32[3,16,4096,221], param_1.752: f32[3,16,128,128], param_2.715: f32[3,4096,221], param_3.568: f32[3,4096,221]) -> f32[4096,221] {
  %param_0.438 = f32[3,16,4096,221]{3,2,1,0} parameter(0)
  %constant_1147 = f32[] constant(1)
  %broadcast.1583 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_1147), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_3.568 = f32[3,4096,221]{2,1,0} parameter(3)
  %floor.207 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_3.568), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.220 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_3.568, f32[3,4096,221]{2,1,0} %floor.207), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.219 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1583, f32[3,4096,221]{2,1,0} %subtract.220), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.715 = f32[3,4096,221]{2,1,0} parameter(2)
  %floor.206 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_2.715), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.218 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_2.715, f32[3,4096,221]{2,1,0} %floor.206), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.217 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1583, f32[3,4096,221]{2,1,0} %subtract.218), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.473 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.219, f32[3,4096,221]{2,1,0} %subtract.217), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1582 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.473), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.752 = f32[3,16,128,128]{3,2,1,0} parameter(1)
  %iota.143 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1135 = s32[] constant(0)
  %broadcast.1581 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1135), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.122 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.207), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1133 = s32[] constant(127)
  %broadcast.1580 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1133), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.111 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1581, s32[3,4096,221]{2,1,0} %convert.122, s32[3,4096,221]{2,1,0} %broadcast.1580), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.237 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.111, s32[3,4096,221]{2,1,0} %broadcast.1581), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1132 = s32[] constant(128)
  %broadcast.1579 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1132), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.681 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.111, s32[3,4096,221]{2,1,0} %broadcast.1579), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.251 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.237, s32[3,4096,221]{2,1,0} %add.681, s32[3,4096,221]{2,1,0} %clamp.111), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.238 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.251), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.121 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.206), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.110 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1581, s32[3,4096,221]{2,1,0} %convert.121, s32[3,4096,221]{2,1,0} %broadcast.1580), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.236 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.110, s32[3,4096,221]{2,1,0} %broadcast.1581), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.680 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.110, s32[3,4096,221]{2,1,0} %broadcast.1579), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.250 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.236, s32[3,4096,221]{2,1,0} %add.680, s32[3,4096,221]{2,1,0} %clamp.110), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.237 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.250), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.182 = s32[3,4096,221,3]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.143, s32[3,4096,221,1]{2,1,0,3} %bitcast.238, s32[3,4096,221,1]{2,1,0,3} %bitcast.237), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.95 = f32[3,16,4096,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_1.752, s32[3,4096,221,3]{2,1,0,3} %concatenate.182), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.472 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.1582, f32[3,16,4096,221]{3,2,1,0} %gather.95), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.471 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.219, f32[3,4096,221]{2,1,0} %subtract.218), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1574 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.471), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.142 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.236 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.251), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1134 = s32[] constant(1)
  %broadcast.1569 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1134), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.678 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.121, s32[3,4096,221]{2,1,0} %broadcast.1569), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.108 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1581, s32[3,4096,221]{2,1,0} %add.678, s32[3,4096,221]{2,1,0} %broadcast.1580), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.234 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.108, s32[3,4096,221]{2,1,0} %broadcast.1581), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.677 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.108, s32[3,4096,221]{2,1,0} %broadcast.1579), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.247 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.234, s32[3,4096,221]{2,1,0} %add.677, s32[3,4096,221]{2,1,0} %clamp.108), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.235 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.247), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.181 = s32[3,4096,221,3]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.142, s32[3,4096,221,1]{3,2,1,0} %bitcast.236, s32[3,4096,221,1]{3,2,1,0} %bitcast.235), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.94 = f32[3,16,4096,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_1.752, s32[3,4096,221,3]{3,2,1,0} %concatenate.181), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.470 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.1574, f32[3,16,4096,221]{3,2,1,0} %gather.94), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.676 = f32[3,16,4096,221]{3,2,1,0} add(f32[3,16,4096,221]{3,2,1,0} %multiply.472, f32[3,16,4096,221]{3,2,1,0} %multiply.470), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.469 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.220, f32[3,4096,221]{2,1,0} %subtract.217), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1565 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.469), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.141 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.675 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.122, s32[3,4096,221]{2,1,0} %broadcast.1569), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.107 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1581, s32[3,4096,221]{2,1,0} %add.675, s32[3,4096,221]{2,1,0} %broadcast.1580), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.233 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.107, s32[3,4096,221]{2,1,0} %broadcast.1581), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.674 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.107, s32[3,4096,221]{2,1,0} %broadcast.1579), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.246 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.233, s32[3,4096,221]{2,1,0} %add.674, s32[3,4096,221]{2,1,0} %clamp.107), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.234 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.246), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.180 = s32[3,4096,221,3]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.141, s32[3,4096,221,1]{2,1,0,3} %bitcast.234, s32[3,4096,221,1]{2,1,0,3} %bitcast.237), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.93 = f32[3,16,4096,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_1.752, s32[3,4096,221,3]{2,1,0,3} %concatenate.180), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.468 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.1565, f32[3,16,4096,221]{3,2,1,0} %gather.93), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.672 = f32[3,16,4096,221]{3,2,1,0} add(f32[3,16,4096,221]{3,2,1,0} %add.676, f32[3,16,4096,221]{3,2,1,0} %multiply.468), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.467 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.220, f32[3,4096,221]{2,1,0} %subtract.218), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1557 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.467), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.140 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.232 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.246), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.179 = s32[3,4096,221,3]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.140, s32[3,4096,221,1]{3,2,1,0} %bitcast.232, s32[3,4096,221,1]{3,2,1,0} %bitcast.235), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.92 = f32[3,16,4096,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_1.752, s32[3,4096,221,3]{3,2,1,0} %concatenate.179), offset_dims={1}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.466 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.1557, f32[3,16,4096,221]{3,2,1,0} %gather.92), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.667 = f32[3,16,4096,221]{3,2,1,0} add(f32[3,16,4096,221]{3,2,1,0} %add.672, f32[3,16,4096,221]{3,2,1,0} %multiply.466), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.378 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %param_0.438, f32[3,16,4096,221]{3,2,1,0} %add.667), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %bitcast.112 = f32[48,4096,221]{2,1,0} bitcast(f32[3,16,4096,221]{3,2,1,0} %multiply.378)
  %constant_648 = f32[] constant(0)
  ROOT %reduce.46 = f32[4096,221]{1,0} reduce(f32[48,4096,221]{2,1,0} %bitcast.112, f32[] %constant_648), dimensions={0}, to_apply=%region_4.507, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
}

%fused_computation.130 (param_0.498: f32[3,4096,221], param_1.672: s32[3,1], param_2.614: s32[3,1]) -> (f32[3,4096,221], f32[3,4096,221], f32[3,4096,221]) {
  %param_0.498 = f32[3,4096,221]{2,1,0} parameter(0)
  %bitcast.158 = f32[1,3,4096,221]{3,2,0,1} bitcast(f32[3,4096,221]{2,1,0} %param_0.498), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 221) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %param_2.614 = s32[3,1]{1,0} parameter(2)
  %gather.79 = f32[3,4096,221]{2,1,0} gather(f32[3,4096,221]{2,1,0} %param_0.498, s32[3,1]{1,0} %param_2.614), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,4096,221}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 4096, 221) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=58}
  %bitcast.157 = f32[1,3,4096,221]{3,2,0,1} bitcast(f32[3,4096,221]{2,1,0} %gather.79), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 221) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %param_1.672 = s32[3,1]{1,0} parameter(1)
  %gather.78 = f32[3,4096,221]{2,1,0} gather(f32[3,4096,221]{2,1,0} %param_0.498, s32[3,1]{1,0} %param_1.672), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,4096,221}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 4096, 221) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=59}
  %bitcast.156 = f32[1,3,4096,221]{3,2,0,1} bitcast(f32[3,4096,221]{2,1,0} %gather.78), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 221) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %concatenate.142 = f32[3,3,4096,221]{3,2,0,1} concatenate(f32[1,3,4096,221]{3,2,0,1} %bitcast.158, f32[1,3,4096,221]{3,2,0,1} %bitcast.157, f32[1,3,4096,221]{3,2,0,1} %bitcast.156), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %constant_742 = f32[] constant(1)
  %broadcast.1179 = f32[3,3,4096,221]{3,2,0,1} broadcast(f32[] %constant_742), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=159}
  %add.514 = f32[3,3,4096,221]{3,2,0,1} add(f32[3,3,4096,221]{3,2,0,1} %concatenate.142, f32[3,3,4096,221]{3,2,0,1} %broadcast.1179), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=159}
  %constant_741 = f32[] constant(63.5)
  %broadcast.1178 = f32[3,3,4096,221]{3,2,0,1} broadcast(f32[] %constant_741), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=162}
  %multiply.421 = f32[3,3,4096,221]{3,2,0,1} multiply(f32[3,3,4096,221]{3,2,0,1} %add.514, f32[3,3,4096,221]{3,2,0,1} %broadcast.1178), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=162}
  %slice.346 = f32[3,1,4096,221]{3,2,0,1} slice(f32[3,3,4096,221]{3,2,0,1} %multiply.421), slice={[0:3], [1:2], [0:4096], [0:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 1, 4096, 221) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.120 = f32[3,4096,221]{2,1,0} bitcast(f32[3,1,4096,221]{3,2,0,1} %slice.346), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.347.clone.1 = f32[3,1,4096,221]{3,2,0,1} slice(f32[3,3,4096,221]{3,2,0,1} %multiply.421), slice={[0:3], [0:1], [0:4096], [0:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1, 2, 3), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(3, 1, 4096, 221) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=165}
  %bitcast.123.clone.1 = f32[3,4096,221]{2,1,0} bitcast(f32[3,1,4096,221]{3,2,0,1} %slice.347.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.345.clone.1 = f32[3,1,4096,221]{3,2,0,1} slice(f32[3,3,4096,221]{3,2,0,1} %multiply.421), slice={[0:3], [2:3], [0:4096], [0:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/slice[start_indices=(0, 1, 0, 0) limit_indices=(3, 2, 4096, 221) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.119.clone.1 = f32[3,4096,221]{2,1,0} bitcast(f32[3,1,4096,221]{3,2,0,1} %slice.345.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %tuple.80 = (f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}) tuple(f32[3,4096,221]{2,1,0} %bitcast.120, f32[3,4096,221]{2,1,0} %bitcast.123.clone.1, f32[3,4096,221]{2,1,0} %bitcast.119.clone.1)
}

%fused_computation.131 (param_0.452: f32[3,16,128], param_1.677: f32[3,4096,221]) -> f32[3,16,4096,221] {
  %constant_765 = f32[] constant(1)
  %broadcast.1201 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_765), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.677 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.57 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.677), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.120 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.677, f32[3,4096,221]{2,1,0} %floor.57), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.119 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1201, f32[3,4096,221]{2,1,0} %subtract.120), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1200 = f32[3,16,4096,221]{2,1,0,3} broadcast(f32[3,4096,221]{2,1,0} %subtract.119), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.452 = f32[3,16,128]{2,1,0} parameter(0)
  %iota.109 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_751 = s32[] constant(0)
  %broadcast.1213 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_751), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.40 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.57), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_749 = s32[] constant(127)
  %broadcast.1212 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_749), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.29 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1213, s32[3,4096,221]{2,1,0} %convert.40, s32[3,4096,221]{2,1,0} %broadcast.1212), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.155 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.29, s32[3,4096,221]{2,1,0} %broadcast.1213), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_748 = s32[] constant(128)
  %broadcast.1211 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_748), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.526 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.29, s32[3,4096,221]{2,1,0} %broadcast.1211), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.167 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.155, s32[3,4096,221]{2,1,0} %add.526, s32[3,4096,221]{2,1,0} %clamp.29), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.164 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.167), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.148 = s32[3,4096,221,2]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.109, s32[3,4096,221,1]{2,1,0,3} %bitcast.164), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.65 = f32[3,16,4096,221]{2,1,0,3} gather(f32[3,16,128]{2,1,0} %param_0.452, s32[3,4096,221,2]{2,1,0,3} %concatenate.148), offset_dims={1}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,16,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 16, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.388 = f32[3,16,4096,221]{2,1,0,3} multiply(f32[3,16,4096,221]{2,1,0,3} %broadcast.1200, f32[3,16,4096,221]{2,1,0,3} %gather.65), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.29 = f32[3,16,4096,221]{3,2,1,0} copy(f32[3,16,4096,221]{2,1,0,3} %multiply.388), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1092 = f32[3,16,4096,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %subtract.120), dimensions={0,2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.105 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_750 = s32[] constant(1)
  %broadcast.1187 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_750), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.518 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.40, s32[3,4096,221]{2,1,0} %broadcast.1187), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.25 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1213, s32[3,4096,221]{2,1,0} %add.518, s32[3,4096,221]{2,1,0} %broadcast.1212), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/jit(vmap(clip))/jit(jit_vmap(clip))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.151 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.25, s32[3,4096,221]{2,1,0} %broadcast.1213), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.517 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.25, s32[3,4096,221]{2,1,0} %broadcast.1211), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.163 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.151, s32[3,4096,221]{2,1,0} %add.517, s32[3,4096,221]{2,1,0} %clamp.25), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.160 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.163), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.144 = s32[3,4096,221,2]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.105, s32[3,4096,221,1]{3,2,1,0} %bitcast.160), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.64 = f32[3,16,4096,221]{3,2,1,0} gather(f32[3,16,128]{2,1,0} %param_0.452, s32[3,4096,221,2]{3,2,1,0} %concatenate.144), offset_dims={1}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,16,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 16, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.387 = f32[3,16,4096,221]{3,2,1,0} multiply(f32[3,16,4096,221]{3,2,1,0} %broadcast.1092, f32[3,16,4096,221]{3,2,1,0} %gather.64), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %add.476 = f32[3,16,4096,221]{3,2,1,0} add(f32[3,16,4096,221]{3,2,1,0} %copy.29, f32[3,16,4096,221]{3,2,1,0} %multiply.387), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%fused_computation.137 (param_0.427: u32[111], param_1.583: u32[111]) -> f32[221] {
  %param_0.427 = u32[111]{0} parameter(0)
  %param_1.583 = u32[111]{0} parameter(1)
  %concatenate.133 = u32[222]{0} concatenate(u32[111]{0} %param_0.427, u32[111]{0} %param_1.583), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %slice.348 = u32[221]{0} slice(u32[222]{0} %concatenate.133), slice={[0:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(), start_index_map=(0,)) slice_sizes=(221,) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_635 = u32[] constant(9)
  %broadcast.1065 = u32[221]{0} broadcast(u32[] %constant_635), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/shift_right_logical" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %shift-right-logical.2 = u32[221]{0} shift-right-logical(u32[221]{0} %slice.348, u32[221]{0} %broadcast.1065), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/shift_right_logical" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_636 = u32[] constant(1065353216)
  %broadcast.1066 = u32[221]{0} broadcast(u32[] %constant_636), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/or" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %or.2 = u32[221]{0} or(u32[221]{0} %shift-right-logical.2, u32[221]{0} %broadcast.1066), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/or" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %bitcast-convert.20 = f32[221]{0} bitcast-convert(u32[221]{0} %or.2), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/bitcast_convert_type[new_dtype=float32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_637 = f32[] constant(-1)
  %broadcast.1068 = f32[221]{0} broadcast(f32[] %constant_637), dimensions={}
  %add.481 = f32[221]{0} add(f32[221]{0} %bitcast-convert.20, f32[221]{0} %broadcast.1068), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_466 = f32[] constant(1.17549435e-38)
  %broadcast.936 = f32[221]{0} broadcast(f32[] %constant_466), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %maximum.62 = f32[221]{0} maximum(f32[221]{0} %add.481, f32[221]{0} %broadcast.936), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %log.6 = f32[221]{0} log(f32[221]{0} %maximum.62), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/log" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %negate.23 = f32[221]{0} negate(f32[221]{0} %log.6), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  ROOT %log.5 = f32[221]{0} log(f32[221]{0} %negate.23), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/log" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
}

%fused_computation.139 (param_0.595: u32[2], param_1.824: u32[2]) -> (u32[111], u32[111], u32[111], u32[111]) {
  %param_0.595 = u32[2]{0} parameter(0)
  %param_1.824 = u32[2]{0} parameter(1)
  %concatenate.194 = u32[4]{0} concatenate(u32[2]{0} %param_0.595, u32[2]{0} %param_1.824), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %bitcast.242 = u32[2,2]{1,0} bitcast(u32[4]{0} %concatenate.194), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(2, 2) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %slice.351 = u32[1,1]{1,0} slice(u32[2,2]{1,0} %bitcast.242), slice={[1:2], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(0,) limit_indices=(1,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %bitcast.129 = u32[] bitcast(u32[1,1]{1,0} %slice.351), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %broadcast.938 = u32[111]{0} broadcast(u32[] %bitcast.129), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %slice.350.clone.1 = u32[1,2]{1,0} slice(u32[2,2]{1,0} %bitcast.242), slice={[1:2], [0:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(1, 0) limit_indices=(2, 2) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %bitcast.128.clone.1 = u32[2]{0} bitcast(u32[1,2]{1,0} %slice.350.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %slice.349.clone.1 = u32[1]{0} slice(u32[2]{0} %bitcast.128.clone.1), slice={[1:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(1,) limit_indices=(2,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %bitcast.127.clone.1 = u32[] bitcast(u32[1]{0} %slice.349.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %broadcast.937.clone.1 = u32[111]{0} broadcast(u32[] %bitcast.127.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %slice.356.clone.1 = u32[1]{0} slice(u32[2]{0} %param_0.595), slice={[1:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(1,) limit_indices=(2,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %bitcast.133.clone.1 = u32[] bitcast(u32[1]{0} %slice.356.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %broadcast.952.clone.1 = u32[111]{0} broadcast(u32[] %bitcast.133.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %slice.357.clone.1 = u32[1]{0} slice(u32[2]{0} %param_0.595), slice={[0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(0,) limit_indices=(1,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %bitcast.134.clone.1 = u32[] bitcast(u32[1]{0} %slice.357.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %broadcast.953.clone.1 = u32[111]{0} broadcast(u32[] %bitcast.134.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  ROOT %tuple.85 = (u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}) tuple(u32[111]{0} %broadcast.938, u32[111]{0} %broadcast.937.clone.1, u32[111]{0} %broadcast.952.clone.1, u32[111]{0} %broadcast.953.clone.1)
}

%fused_computation.141 (param_0.464: f32[2,3], param_1.642: u32[111], param_2.588: u32[111], param_3.449: f32[4096,3], param_4.248: f32[4096,3], param_5.147: pred[4096], param_6.93: f32[4096], param_7.42: f32[4096]) -> f32[3,4096,221] {
  %param_4.248 = f32[4096,3]{1,0} parameter(4)
  %broadcast.948 = f32[4096,3,221]{2,0,1} broadcast(f32[4096,3]{1,0} %param_4.248), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=335}
  %param_3.449 = f32[4096,3]{1,0} parameter(3)
  %broadcast.947 = f32[4096,3,221]{2,0,1} broadcast(f32[4096,3]{1,0} %param_3.449), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=335}
  %param_5.147 = pred[4096]{0} parameter(5)
  %param_6.93 = f32[4096]{0} parameter(6)
  %constant_679 = f32[] constant(0)
  %broadcast.1131 = f32[4096]{0} broadcast(f32[] %constant_679), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.147 = f32[4096]{0} select(pred[4096]{0} %param_5.147, f32[4096]{0} %param_6.93, f32[4096]{0} %broadcast.1131), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %broadcast.946 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %select.147), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=331}
  %iota.99 = s32[221]{0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/iota[dtype=int32 shape=(221,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=322}
  %convert.20 = f32[221]{0} convert(s32[221]{0} %iota.99), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %param_1.642 = u32[111]{0} parameter(1)
  %param_2.588 = u32[111]{0} parameter(2)
  %concatenate.135 = u32[222]{0} concatenate(u32[111]{0} %param_1.642, u32[111]{0} %param_2.588), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %slice.354 = u32[221]{0} slice(u32[222]{0} %concatenate.135), slice={[0:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(), start_index_map=(0,)) slice_sizes=(221,) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %constant_676 = u32[] constant(9)
  %broadcast.1103 = u32[221]{0} broadcast(u32[] %constant_676), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/shift_right_logical" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %shift-right-logical.3 = u32[221]{0} shift-right-logical(u32[221]{0} %slice.354, u32[221]{0} %broadcast.1103), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/shift_right_logical" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %constant_677 = u32[] constant(1065353216)
  %broadcast.1104 = u32[221]{0} broadcast(u32[] %constant_677), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/or" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %or.3 = u32[221]{0} or(u32[221]{0} %shift-right-logical.3, u32[221]{0} %broadcast.1104), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/or" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %bitcast-convert.21 = f32[221]{0} bitcast-convert(u32[221]{0} %or.3), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/bitcast_convert_type[new_dtype=float32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %constant_678 = f32[] constant(-1)
  %broadcast.1105 = f32[221]{0} broadcast(f32[] %constant_678), dimensions={}
  %add.487 = f32[221]{0} add(f32[221]{0} %bitcast-convert.21, f32[221]{0} %broadcast.1105), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %broadcast.945 = f32[221]{0} broadcast(f32[] %constant_679), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %maximum.63 = f32[221]{0} maximum(f32[221]{0} %add.487, f32[221]{0} %broadcast.945), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %add.486 = f32[221]{0} add(f32[221]{0} %convert.20, f32[221]{0} %maximum.63), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %broadcast.944 = f32[4096,221]{1,0} broadcast(f32[221]{0} %add.486), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %param_7.42 = f32[4096]{0} parameter(7)
  %constant_713 = f32[] constant(0.001)
  %broadcast.1149 = f32[4096]{0} broadcast(f32[] %constant_713), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.503 = f32[4096]{0} add(f32[4096]{0} %param_6.93, f32[4096]{0} %broadcast.1149), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.69 = f32[4096]{0} maximum(f32[4096]{0} %param_7.42, f32[4096]{0} %add.503), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.153 = f32[4096]{0} select(pred[4096]{0} %param_5.147, f32[4096]{0} %maximum.69, f32[4096]{0} %broadcast.1149), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %subtract.112 = f32[4096]{0} subtract(f32[4096]{0} %select.153, f32[4096]{0} %select.147), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_711 = f32[] constant(0.00452488707)
  %broadcast.1145 = f32[4096]{0} broadcast(f32[] %constant_711), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.410 = f32[4096]{0} multiply(f32[4096]{0} %subtract.112, f32[4096]{0} %broadcast.1145), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.1144 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.410), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.393 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %broadcast.944, f32[4096,221]{1,0} %broadcast.1144), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %add.485 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %broadcast.946, f32[4096,221]{1,0} %multiply.393), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=331}
  %broadcast.943 = f32[4096,3,221]{2,0,1} broadcast(f32[4096,221]{1,0} %add.485), dimensions={0,2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=335}
  %multiply.392 = f32[4096,3,221]{2,0,1} multiply(f32[4096,3,221]{2,0,1} %broadcast.947, f32[4096,3,221]{2,0,1} %broadcast.943), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=335}
  %add.484 = f32[4096,3,221]{2,0,1} add(f32[4096,3,221]{2,0,1} %broadcast.948, f32[4096,3,221]{2,0,1} %multiply.392), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=335}
  %transpose.12 = f32[3,4096,221]{2,1,0} transpose(f32[4096,3,221]{2,0,1} %add.484), dimensions={1,0,2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/transpose[permutation=(1, 0, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=114}
  %param_0.464 = f32[2,3]{1,0} parameter(0)
  %slice.353 = f32[1,3]{1,0} slice(f32[2,3]{1,0} %param_0.464), slice={[0:1], [0:3]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 3) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %bitcast.132 = f32[3]{0} bitcast(f32[1,3]{1,0} %slice.353), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %broadcast.942 = f32[3,4096,221]{2,1,0} broadcast(f32[3]{0} %bitcast.132), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %subtract.74 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %transpose.12, f32[3,4096,221]{2,1,0} %broadcast.942), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %slice.352 = f32[1,3]{1,0} slice(f32[2,3]{1,0} %param_0.464), slice={[1:2], [0:3]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 3) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %subtract.73 = f32[1,3]{1,0} subtract(f32[1,3]{1,0} %slice.352, f32[1,3]{1,0} %slice.353), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %bitcast.131 = f32[3]{0} bitcast(f32[1,3]{1,0} %subtract.73), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %broadcast.941 = f32[3,4096,221]{2,1,0} broadcast(f32[3]{0} %bitcast.131), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %divide.48 = f32[3,4096,221]{2,1,0} divide(f32[3,4096,221]{2,1,0} %subtract.74, f32[3,4096,221]{2,1,0} %broadcast.941), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %constant_468 = f32[] constant(-0.5)
  %broadcast.940 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_468), dimensions={}
  %add.482 = f32[3,4096,221]{2,1,0} add(f32[3,4096,221]{2,1,0} %divide.48, f32[3,4096,221]{2,1,0} %broadcast.940), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %constant_680 = f32[] constant(2)
  %broadcast.939 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_680), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=128}
  ROOT %multiply.391 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %add.482, f32[3,4096,221]{2,1,0} %broadcast.939), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=128}
}

%fused_computation.143 () -> u32[111] {
  %iota.100 = u32[221]{0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/iota[dtype=uint32 shape=(221,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %constant_470 = u32[1]{0} constant({0})
  %concatenate.136 = u32[222]{0} concatenate(u32[221]{0} %iota.100, u32[1]{0} %constant_470), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  ROOT %slice.355 = u32[111]{0} slice(u32[222]{0} %concatenate.136), slice={[111:222]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(111,) limit_indices=(222,) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
}

%fused_computation.149 (param_0.358: u32[2]) -> (u32[2], u32[2]) {
  %param_0.358 = u32[2]{0} parameter(0)
  %slice.363 = u32[1]{0} slice(u32[2]{0} %param_0.358), slice={[0:1]}, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(0,) limit_indices=(1,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %bitcast.138 = u32[] bitcast(u32[1]{0} %slice.363), metadata={op_name="jit(training_step)/jit(main)/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %broadcast.957 = u32[2]{0} broadcast(u32[] %bitcast.138), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %slice.362.clone.1 = u32[1]{0} slice(u32[2]{0} %param_0.358), slice={[1:2]}, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(1,) limit_indices=(2,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %bitcast.137.clone.1 = u32[] bitcast(u32[1]{0} %slice.362.clone.1), metadata={op_name="jit(training_step)/jit(main)/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %broadcast.956.clone.1 = u32[2]{0} broadcast(u32[] %bitcast.137.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  ROOT %tuple.88 = (u32[2]{0}, u32[2]{0}) tuple(u32[2]{0} %broadcast.957, u32[2]{0} %broadcast.956.clone.1)
}

%fused_computation.152 (param_0.395: f32[135168,128]) -> f32[135168,128] {
  %param_0.395 = f32[135168,128]{1,0} parameter(0)
  %constant_598 = f32[] constant(0)
  %broadcast.1024 = f32[135168,128]{1,0} broadcast(f32[] %constant_598), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gt" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  ROOT %maximum.66 = f32[135168,128]{1,0} maximum(f32[135168,128]{1,0} %param_0.395, f32[135168,128]{1,0} %broadcast.1024), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(relu)/jit(jit_relu)/max" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=105}
}

%fused_computation.153 (param_0.396: f32[135168,128]) -> f32[135168,128] {
  %param_0.396 = f32[135168,128]{1,0} parameter(0)
  %constant_599 = f32[] constant(0)
  %broadcast.1026 = f32[135168,128]{1,0} broadcast(f32[] %constant_599), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gt" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  ROOT %maximum.67 = f32[135168,128]{1,0} maximum(f32[135168,128]{1,0} %param_0.396, f32[135168,128]{1,0} %broadcast.1026), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(relu)/jit(jit_relu)/max" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
}

%fused_computation.154 () -> u32[111] {
  %iota.101 = u32[221]{0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/iota[dtype=uint32 shape=(221,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  ROOT %slice.366 = u32[111]{0} slice(u32[221]{0} %iota.101), slice={[0:111]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(0,) limit_indices=(111,) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
}

%fused_computation.155 () -> u32[2] {
  %iota.102 = u32[4]{0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/iota[dtype=uint32 shape=(4,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  ROOT %slice.367 = u32[2]{0} slice(u32[4]{0} %iota.102), slice={[2:4]}, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(2,) limit_indices=(4,) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
}

%fused_computation.156 () -> u32[2] {
  %iota.103 = u32[4]{0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/iota[dtype=uint32 shape=(4,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  ROOT %slice.368 = u32[2]{0} slice(u32[4]{0} %iota.103), slice={[0:2]}, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(0,) limit_indices=(2,) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
}

%region_1.225 (Arg_0.226: f32[], Arg_1.227: f32[]) -> f32[] {
  %Arg_0.226 = f32[] parameter(0)
  %Arg_1.227 = f32[] parameter(1)
  ROOT %maximum.228 = f32[] maximum(f32[] %Arg_0.226, f32[] %Arg_1.227), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_max[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=368}
}

%region_0.221 (Arg_0.222: f32[], Arg_1.223: f32[]) -> f32[] {
  %Arg_0.222 = f32[] parameter(0)
  %Arg_1.223 = f32[] parameter(1)
  ROOT %minimum.224 = f32[] minimum(f32[] %Arg_0.222, f32[] %Arg_1.223), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_min[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=367}
}

%fused_computation.159 (param_0.750: f32[4096,3], param_1.1004: f32[4096,3], param_2.1050: f32[2,3]) -> (f32[4096,3], f32[4096,3]) {
  %param_2.1050 = f32[2,3]{1,0} parameter(2)
  %broadcast.1129 = f32[4096,2,3]{2,1,0} broadcast(f32[2,3]{1,0} %param_2.1050), dimensions={1,2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=361}
  %param_1.1004 = f32[4096,3]{1,0} parameter(1)
  %broadcast.1128 = f32[4096,2,3]{2,1,0} broadcast(f32[4096,3]{1,0} %param_1.1004), dimensions={0,2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=361}
  %subtract.110 = f32[4096,2,3]{2,1,0} subtract(f32[4096,2,3]{2,1,0} %broadcast.1129, f32[4096,2,3]{2,1,0} %broadcast.1128), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=361}
  %param_0.750 = f32[4096,3]{1,0} parameter(0)
  %constant_700 = f32[] constant(1e-08)
  %broadcast.1126 = f32[4096,3]{1,0} broadcast(f32[] %constant_700), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=363}
  %add.501 = f32[4096,3]{1,0} add(f32[4096,3]{1,0} %param_0.750, f32[4096,3]{1,0} %broadcast.1126), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=363}
  %broadcast.1124 = f32[4096,2,3]{2,1,0} broadcast(f32[4096,3]{1,0} %add.501), dimensions={0,2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=362}
  %divide.53 = f32[4096,2,3]{2,1,0} divide(f32[4096,2,3]{2,1,0} %subtract.110, f32[4096,2,3]{2,1,0} %broadcast.1124), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=362}
  %constant_1871 = f32[] constant(-inf)
  %reduce.49 = f32[4096,3]{1,0} reduce(f32[4096,2,3]{2,1,0} %divide.53, f32[] %constant_1871), dimensions={1}, to_apply=%region_1.225, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_max[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=368}
  %constant_1879 = f32[] constant(inf)
  %reduce.48.clone.1 = f32[4096,3]{1,0} reduce(f32[4096,2,3]{2,1,0} %divide.53, f32[] %constant_1879), dimensions={1}, to_apply=%region_0.221, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_min[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=367}
  ROOT %tuple.83 = (f32[4096,3]{1,0}, f32[4096,3]{1,0}) tuple(f32[4096,3]{1,0} %reduce.49, f32[4096,3]{1,0} %reduce.48.clone.1)
}

%fused_computation.161 (param_0.755: f32[4096,13], param_1.1008: f32[4096,6], param_2.1054: f32[4096,6], param_3.930: f32[4096,55]) -> (f32[4096,27], f32[4096,27]) {
  %param_0.755 = f32[4096,13]{1,0} parameter(0)
  %slice.400 = f32[4096,1]{1,0} slice(f32[4096,13]{1,0} %param_0.755), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_1.1008 = f32[4096,6]{1,0} parameter(1)
  %param_2.1054 = f32[4096,6]{1,0} parameter(2)
  %add.729 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %param_1.1008, f32[4096,6]{1,0} %param_2.1054), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.399 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_0.755), slice={[0:4096:1], [2:13:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 13) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.728 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %add.729, f32[4096,6]{1,0} %slice.399), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.186 = f32[4096,7]{1,0} concatenate(f32[4096,1]{1,0} %slice.400, f32[4096,6]{1,0} %add.728), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_1238 = f32[] constant(0)
  %pad.133 = f32[4096,13]{1,0} pad(f32[4096,7]{1,0} %concatenate.186, f32[] %constant_1238), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.132 = f32[4096,13]{1,0} pad(f32[4096,6]{1,0} %add.729, f32[] %constant_1238), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.727 = f32[4096,13]{1,0} add(f32[4096,13]{1,0} %pad.133, f32[4096,13]{1,0} %pad.132), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.129 = f32[4096,27]{1,0} pad(f32[4096,13]{1,0} %add.727, f32[] %constant_1238), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_3.930 = f32[4096,55]{1,0} parameter(3)
  %slice.380.clone.1 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_3.930), slice={[0:4096:1], [0:54:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 54) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.379.clone.1 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_3.930), slice={[0:4096:1], [1:55:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 55) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.709.clone.1 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %slice.380.clone.1, f32[4096,27]{1,0} %slice.379.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.323.clone.1 = f32[4096,1]{1,0} slice(f32[4096,27]{1,0} %add.709.clone.1), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.322.clone.1 = f32[4096,13]{1,0} slice(f32[4096,27]{1,0} %add.709.clone.1), slice={[0:4096:1], [2:27:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 27) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.446.clone.1 = f32[4096,13]{1,0} add(f32[4096,13]{1,0} %add.727, f32[4096,13]{1,0} %slice.322.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.122.clone.1 = f32[4096,14]{1,0} concatenate(f32[4096,1]{1,0} %slice.323.clone.1, f32[4096,13]{1,0} %add.446.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.119.clone.1 = f32[4096,27]{1,0} pad(f32[4096,14]{1,0} %concatenate.122.clone.1, f32[] %constant_1238), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %tuple.78 = (f32[4096,27]{1,0}, f32[4096,27]{1,0}) tuple(f32[4096,27]{1,0} %pad.129, f32[4096,27]{1,0} %pad.119.clone.1)
}

%fused_computation.162 (param_0.754: f32[4096,55], param_1.1007: f32[4096,27], param_2.1053: f32[4096,27], param_3.929: pred[4096], param_4.714: f32[4096], param_5.683: f32[4096], param_6.646: f32[4096,221]) -> (f32[4096,110], f32[4096,110]) {
  %param_0.754 = f32[4096,55]{1,0} parameter(0)
  %slice.408 = f32[4096,1]{1,0} slice(f32[4096,55]{1,0} %param_0.754), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_1.1007 = f32[4096,27]{1,0} parameter(1)
  %param_2.1053 = f32[4096,27]{1,0} parameter(2)
  %add.741 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %param_1.1007, f32[4096,27]{1,0} %param_2.1053), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.407 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_0.754), slice={[0:4096:1], [2:55:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 55) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.740 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %add.741, f32[4096,27]{1,0} %slice.407), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.190 = f32[4096,28]{1,0} concatenate(f32[4096,1]{1,0} %slice.408, f32[4096,27]{1,0} %add.740), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_1244 = f32[] constant(0)
  %pad.142 = f32[4096,55]{1,0} pad(f32[4096,28]{1,0} %concatenate.190, f32[] %constant_1244), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.141 = f32[4096,55]{1,0} pad(f32[4096,27]{1,0} %add.741, f32[] %constant_1244), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.739 = f32[4096,55]{1,0} add(f32[4096,55]{1,0} %pad.142, f32[4096,55]{1,0} %pad.141), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.138 = f32[4096,110]{1,0} pad(f32[4096,55]{1,0} %add.739, f32[] %constant_1244), padding=0_0_0x1_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_6.646 = f32[4096,221]{1,0} parameter(6)
  %negate.37.clone.1 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_6.646), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_3.929 = pred[4096]{0} parameter(3)
  %param_5.683 = f32[4096]{0} parameter(5)
  %param_4.714 = f32[4096]{0} parameter(4)
  %constant_1221_clone_1 = f32[] constant(0.001)
  %broadcast.1656.clone.1 = f32[4096]{0} broadcast(f32[] %constant_1221_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.703.clone.1 = f32[4096]{0} add(f32[4096]{0} %param_4.714, f32[4096]{0} %broadcast.1656.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.87.clone.1 = f32[4096]{0} maximum(f32[4096]{0} %param_5.683, f32[4096]{0} %add.703.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.279.clone.1 = f32[4096]{0} select(pred[4096]{0} %param_3.929, f32[4096]{0} %maximum.87.clone.1, f32[4096]{0} %broadcast.1656.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %broadcast.1655.clone.1 = f32[4096]{0} broadcast(f32[] %constant_1244), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.278.clone.1 = f32[4096]{0} select(pred[4096]{0} %param_3.929, f32[4096]{0} %param_4.714, f32[4096]{0} %broadcast.1655.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.238.clone.1 = f32[4096]{0} subtract(f32[4096]{0} %select.279.clone.1, f32[4096]{0} %select.278.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1219_clone_1 = f32[] constant(0.00452488707)
  %broadcast.1654.clone.1 = f32[4096]{0} broadcast(f32[] %constant_1219_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.503.clone.1 = f32[4096]{0} multiply(f32[4096]{0} %subtract.238.clone.1, f32[4096]{0} %broadcast.1654.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.1653.clone.1 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.503.clone.1), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.502.clone.1 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.37.clone.1, f32[4096,221]{1,0} %broadcast.1653.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %slice.372.clone.1 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %multiply.502.clone.1), slice={[0:4096:1], [0:220:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 220) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.371.clone.1 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %multiply.502.clone.1), slice={[0:4096:1], [1:221:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.702.clone.1 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %slice.372.clone.1, f32[4096,110]{1,0} %slice.371.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.319.clone.1 = f32[4096,1]{1,0} slice(f32[4096,110]{1,0} %add.702.clone.1), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.318.clone.1 = f32[4096,54]{1,0} slice(f32[4096,55]{1,0} %add.739), slice={[0:4096], [0:54]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 54) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.317.clone.1 = f32[4096,54]{1,0} slice(f32[4096,110]{1,0} %add.702.clone.1), slice={[0:4096:1], [2:110:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 110) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.442.clone.1 = f32[4096,54]{1,0} add(f32[4096,54]{1,0} %slice.318.clone.1, f32[4096,54]{1,0} %slice.317.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.120.clone.1 = f32[4096,55]{1,0} concatenate(f32[4096,1]{1,0} %slice.319.clone.1, f32[4096,54]{1,0} %add.442.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.116.clone.1 = f32[4096,110]{1,0} pad(f32[4096,55]{1,0} %concatenate.120.clone.1, f32[] %constant_1244), padding=0_0_0x0_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %tuple.77 = (f32[4096,110]{1,0}, f32[4096,110]{1,0}) tuple(f32[4096,110]{1,0} %pad.138, f32[4096,110]{1,0} %pad.116.clone.1)
}

%fused_computation.164 (param_0.753: f32[4096,13], param_1.1006: f32[4096,6], param_2.1052: f32[4096,6], param_3.928: f32[4096,55]) -> (f32[4096,27], f32[4096,27]) {
  %param_0.753 = f32[4096,13]{1,0} parameter(0)
  %slice.686 = f32[4096,1]{1,0} slice(f32[4096,13]{1,0} %param_0.753), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_1.1006 = f32[4096,6]{1,0} parameter(1)
  %param_2.1052 = f32[4096,6]{1,0} parameter(2)
  %add.967 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %param_1.1006, f32[4096,6]{1,0} %param_2.1052), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.685 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_0.753), slice={[0:4096:1], [2:13:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 13) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.966 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %add.967, f32[4096,6]{1,0} %slice.685), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.247 = f32[4096,7]{1,0} concatenate(f32[4096,1]{1,0} %slice.686, f32[4096,6]{1,0} %add.966), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_1735 = f32[] constant(0)
  %pad.155 = f32[4096,13]{1,0} pad(f32[4096,7]{1,0} %concatenate.247, f32[] %constant_1735), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.154 = f32[4096,13]{1,0} pad(f32[4096,6]{1,0} %add.967, f32[] %constant_1735), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.965 = f32[4096,13]{1,0} add(f32[4096,13]{1,0} %pad.155, f32[4096,13]{1,0} %pad.154), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.151 = f32[4096,27]{1,0} pad(f32[4096,13]{1,0} %add.965, f32[] %constant_1735), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_3.928 = f32[4096,55]{1,0} parameter(3)
  %slice.666.clone.1 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_3.928), slice={[0:4096:1], [0:54:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 54) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.665.clone.1 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_3.928), slice={[0:4096:1], [1:55:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 55) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.947.clone.1 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %slice.666.clone.1, f32[4096,27]{1,0} %slice.665.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.283.clone.1 = f32[4096,1]{1,0} slice(f32[4096,27]{1,0} %add.947.clone.1), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.282.clone.1 = f32[4096,13]{1,0} slice(f32[4096,27]{1,0} %add.947.clone.1), slice={[0:4096:1], [2:27:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 27) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.350.clone.1 = f32[4096,13]{1,0} add(f32[4096,13]{1,0} %add.965, f32[4096,13]{1,0} %slice.282.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.104.clone.1 = f32[4096,14]{1,0} concatenate(f32[4096,1]{1,0} %slice.283.clone.1, f32[4096,13]{1,0} %add.350.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.101.clone.1 = f32[4096,27]{1,0} pad(f32[4096,14]{1,0} %concatenate.104.clone.1, f32[] %constant_1735), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %tuple.50 = (f32[4096,27]{1,0}, f32[4096,27]{1,0}) tuple(f32[4096,27]{1,0} %pad.151, f32[4096,27]{1,0} %pad.101.clone.1)
}

%fused_computation.165 (param_0.752: f32[4096,55], param_1.1005: f32[4096,27], param_2.1051: f32[4096,27], param_3.927: f32[4096,221]) -> (f32[4096,110], f32[4096,110]) {
  %param_0.752 = f32[4096,55]{1,0} parameter(0)
  %slice.694 = f32[4096,1]{1,0} slice(f32[4096,55]{1,0} %param_0.752), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_1.1005 = f32[4096,27]{1,0} parameter(1)
  %param_2.1051 = f32[4096,27]{1,0} parameter(2)
  %add.979 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %param_1.1005, f32[4096,27]{1,0} %param_2.1051), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.693 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_0.752), slice={[0:4096:1], [2:55:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 55) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.978 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %add.979, f32[4096,27]{1,0} %slice.693), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.251 = f32[4096,28]{1,0} concatenate(f32[4096,1]{1,0} %slice.694, f32[4096,27]{1,0} %add.978), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_1741 = f32[] constant(0)
  %pad.164 = f32[4096,55]{1,0} pad(f32[4096,28]{1,0} %concatenate.251, f32[] %constant_1741), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.163 = f32[4096,55]{1,0} pad(f32[4096,27]{1,0} %add.979, f32[] %constant_1741), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.977 = f32[4096,55]{1,0} add(f32[4096,55]{1,0} %pad.164, f32[4096,55]{1,0} %pad.163), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.160 = f32[4096,110]{1,0} pad(f32[4096,55]{1,0} %add.977, f32[] %constant_1741), padding=0_0_0x1_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_3.927 = f32[4096,221]{1,0} parameter(3)
  %slice.658.clone.1 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_3.927), slice={[0:4096:1], [2:221:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 220) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.17.clone.1 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.658.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 220) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.657.clone.1 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_3.927), slice={[0:4096:1], [1:220:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.16.clone.1 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.657.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.942.clone.1 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %reverse.17.clone.1, f32[4096,110]{1,0} %reverse.16.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.279.clone.1 = f32[4096,1]{1,0} slice(f32[4096,110]{1,0} %add.942.clone.1), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.278.clone.1 = f32[4096,54]{1,0} slice(f32[4096,55]{1,0} %add.977), slice={[0:4096], [0:54]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 54) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.277.clone.1 = f32[4096,54]{1,0} slice(f32[4096,110]{1,0} %add.942.clone.1), slice={[0:4096:1], [2:110:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 110) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.346.clone.1 = f32[4096,54]{1,0} add(f32[4096,54]{1,0} %slice.278.clone.1, f32[4096,54]{1,0} %slice.277.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.102.clone.1 = f32[4096,55]{1,0} concatenate(f32[4096,1]{1,0} %slice.279.clone.1, f32[4096,54]{1,0} %add.346.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.98.clone.1 = f32[4096,110]{1,0} pad(f32[4096,55]{1,0} %concatenate.102.clone.1, f32[] %constant_1741), padding=0_0_0x0_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %tuple.49 = (f32[4096,110]{1,0}, f32[4096,110]{1,0}) tuple(f32[4096,110]{1,0} %pad.160, f32[4096,110]{1,0} %pad.98.clone.1)
}

%fused_computation.166 (param_0.751: s32[]) -> (f32[], f32[], s32[]) {
  %param_0.751 = s32[] parameter(0)
  %constant_692_clone_1 = s32[] constant(0)
  %compare.131.clone.1 = pred[] compare(s32[] %param_0.751, s32[] %constant_692_clone_1), direction=LE, metadata={op_name="jit(training_step)/jit(main)/le" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=199}
  %constant_693_clone_1 = f32[] constant(1)
  %constant_694_clone_1 = f32[] constant(0.1)
  %convert.19.clone.1 = f32[] convert(s32[] %param_0.751), metadata={op_name="jit(training_step)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=195}
  %constant_455_clone_1 = f32[] constant(3.33333337e-05)
  %multiply.346.clone.1 = f32[] multiply(f32[] %convert.19.clone.1, f32[] %constant_455_clone_1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=195}
  %power.2.clone.1 = f32[] power(f32[] %constant_694_clone_1, f32[] %multiply.346.clone.1), metadata={op_name="jit(training_step)/jit(main)/pow" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=199}
  %select.124.clone.1 = f32[] select(pred[] %compare.131.clone.1, f32[] %constant_693_clone_1, f32[] %power.2.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(_where)/jit(jit__where)/select_n" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=198}
  %maximum.60.clone.1 = f32[] maximum(f32[] %select.124.clone.1, f32[] %constant_694_clone_1), metadata={op_name="jit(training_step)/jit(main)/max" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=201}
  %constant_1873 = f32[] constant(0.001)
  %multiply.675 = f32[] multiply(f32[] %maximum.60.clone.1, f32[] %constant_1873), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=356}
  %constant_1872 = s32[] constant(1)
  %add.1064 = s32[] add(s32[] %param_0.751, s32[] %constant_1872), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=364}
  ROOT %tuple.69 = (f32[], f32[], s32[]) tuple(f32[] %multiply.675, f32[] %maximum.60.clone.1, s32[] %add.1064)
}

%region_3.233 (Arg_0.234: f32[], Arg_1.235: f32[]) -> f32[] {
  %Arg_0.234 = f32[] parameter(0)
  %Arg_1.235 = f32[] parameter(1)
  ROOT %minimum.236 = f32[] minimum(f32[] %Arg_0.234, f32[] %Arg_1.235), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_min[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=373}
}

%region_2.229 (Arg_0.230: f32[], Arg_1.231: f32[]) -> f32[] {
  %Arg_0.230 = f32[] parameter(0)
  %Arg_1.231 = f32[] parameter(1)
  ROOT %maximum.232 = f32[] maximum(f32[] %Arg_0.230, f32[] %Arg_1.231), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_max[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
}

%fused_computation.167 (param_0.749: f32[4096,3], param_1.1009: f32[4096,3]) -> (pred[4096], f32[4096], f32[4096]) {
  %param_1.1009 = f32[4096,3]{1,0} parameter(1)
  %constant_688_clone_1 = f32[] constant(-inf)
  %reduce.47.clone.1 = f32[4096]{0} reduce(f32[4096,3]{1,0} %param_1.1009, f32[] %constant_688_clone_1), dimensions={1}, to_apply=%region_2.229, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_max[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %constant_690_clone_1 = f32[] constant(0)
  %broadcast.1107.clone.1 = f32[4096]{0} broadcast(f32[] %constant_690_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %maximum.65.clone.1 = f32[4096]{0} maximum(f32[4096]{0} %reduce.47.clone.1, f32[4096]{0} %broadcast.1107.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %param_0.749 = f32[4096,3]{1,0} parameter(0)
  %constant_1878 = f32[] constant(inf)
  %reduce.50 = f32[4096]{0} reduce(f32[4096,3]{1,0} %param_0.749, f32[] %constant_1878), dimensions={1}, to_apply=%region_3.233, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_min[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=373}
  %compare.360 = pred[4096]{0} compare(f32[4096]{0} %maximum.65.clone.1, f32[4096]{0} %reduce.50), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=375}
  ROOT %tuple.82 = (pred[4096]{0}, f32[4096]{0}, f32[4096]{0}) tuple(pred[4096]{0} %compare.360, f32[4096]{0} %reduce.50, f32[4096]{0} %maximum.65.clone.1)
}

%input_fused_computation_reduce.7 (param_0.759: f32[368,128], param_1.1014: f32[], param_2.1058: f32[368,128]) -> (f32[128], f32[128]) {
  %param_0.759 = f32[368,128]{1,0} parameter(0)
  %param_1.1014 = f32[] parameter(1)
  %reduce.51 = f32[128]{0} reduce(f32[368,128]{1,0} %param_0.759, f32[] %param_1.1014), dimensions={0}, to_apply=%region_28.1714, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  %param_2.1058 = f32[368,128]{1,0} parameter(2)
  %reduce.52.clone.1 = f32[128]{0} reduce(f32[368,128]{1,0} %param_2.1058, f32[] %param_1.1014), dimensions={0}, to_apply=%region_29.1718, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  ROOT %tuple.89 = (f32[128]{0}, f32[128]{0}) tuple(f32[128]{0} %reduce.51, f32[128]{0} %reduce.52.clone.1)
}

ENTRY %main.2261 (Arg_0.1: f32[144,27], Arg_1.2: f32[128], Arg_2.3: f32[150,128], Arg_3.4: f32[128], Arg_4.5: f32[128,128], Arg_5.6: f32[3], Arg_6.7: f32[128,3], Arg_7.8: f32[3,48,128], Arg_8.9: f32[3,48,128,128], Arg_9.10: f32[3,16,128], Arg_10.11: f32[3,16,128,128], Arg_11.12: s32[], Arg_12.13: f32[144,27], Arg_13.14: f32[128], Arg_14.15: f32[150,128], Arg_15.16: f32[128], Arg_16.17: f32[128,128], Arg_17.18: f32[3], Arg_18.19: f32[128,3], Arg_19.20: f32[3,48,128], Arg_20.21: f32[3,48,128,128], Arg_21.22: f32[3,16,128], Arg_22.23: f32[3,16,128,128], Arg_23.24: f32[144,27], Arg_24.25: f32[128], Arg_25.26: f32[150,128], Arg_26.27: f32[128], Arg_27.28: f32[128,128], Arg_28.29: f32[3], Arg_29.30: f32[128,3], Arg_30.31: f32[3,48,128], Arg_31.32: f32[3,48,128,128], Arg_32.33: f32[3,16,128], Arg_33.34: f32[3,16,128,128], Arg_34.35: f32[2,3], Arg_35.36: u32[2], Arg_36.37: s32[], Arg_37.38: f32[4096,3], Arg_38.39: f32[4096,3], Arg_39.40: f32[4096,3]) -> (f32[144,27], f32[128], f32[150,128], f32[128], f32[128,128], /*index=5*/f32[3], f32[128,3], f32[3,48,128], f32[3,48,128,128], f32[3,16,128], /*index=10*/f32[3,16,128,128], s32[], f32[144,27], f32[128], f32[150,128], /*index=15*/f32[128], f32[128,128], f32[3], f32[128,3], f32[3,48,128], /*index=20*/f32[3,48,128,128], f32[3,16,128], f32[3,16,128,128], f32[144,27], f32[128], /*index=25*/f32[150,128], f32[128], f32[128,128], f32[3], f32[128,3], /*index=30*/f32[3,48,128], f32[3,48,128,128], f32[3,16,128], f32[3,16,128,128], f32[2,3], /*index=35*/u32[2], s32[], f32[], f32[], f32[], /*index=40*/f32[]) {
  %Arg_0.1 = f32[144,27]{1,0} parameter(0)
  %copy.43 = f32[144,27]{1,0} copy(f32[144,27]{1,0} %Arg_0.1)
  %Arg_11.12 = s32[] parameter(11)
  %fusion.56 = (f32[], f32[], s32[]) fusion(s32[] %Arg_11.12), kind=kLoop, calls=%fused_computation.56, metadata={op_name="jit(training_step)/jit(main)/sub" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=86}
  %get-tuple-element.186 = f32[] get-tuple-element((f32[], f32[], s32[]) %fusion.56), index=0
  %get-tuple-element.187 = f32[] get-tuple-element((f32[], f32[], s32[]) %fusion.56), index=1
  %Arg_36.37 = s32[] parameter(36), metadata={op_name="jit(training_step)/jit(main)/sub" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=194}
  %fusion.166 = (f32[], f32[], s32[]) fusion(s32[] %Arg_36.37), kind=kLoop, calls=%fused_computation.166, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=356}
  %get-tuple-element.184 = f32[] get-tuple-element((f32[], f32[], s32[]) %fusion.166), index=1
  %Arg_23.24 = f32[144,27]{1,0} parameter(23)
  %Arg_7.8 = f32[3,48,128]{2,1,0} parameter(7)
  %copy.50 = f32[3,48,128]{2,1,0} copy(f32[3,48,128]{2,1,0} %Arg_7.8)
  %constant_367 = s32[3,1]{1,0} constant({ {1}, {2}, {0} }), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(3, 1) broadcast_dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=59}
  %constant_363 = s32[3,1]{1,0} constant({ {2}, {0}, {1} }), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(3, 1) broadcast_dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=58}
  %Arg_34.35 = f32[2,3]{1,0} parameter(34)
  %Arg_35.36 = u32[2]{0} parameter(35)
  %fusion.149 = (u32[2]{0}, u32[2]{0}) fusion(u32[2]{0} %Arg_35.36), kind=kLoop, calls=%fused_computation.149, metadata={op_name="jit(training_step)/jit(main)/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %get-tuple-element.216 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}) %fusion.149), index=0
  %get-tuple-element.217 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}) %fusion.149), index=1
  %fusion.156 = u32[2]{0} fusion(), kind=kLoop, calls=%fused_computation.156, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(0,) limit_indices=(2,) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %fusion.155 = u32[2]{0} fusion(), kind=kLoop, calls=%fused_computation.155, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(2,) limit_indices=(4,) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %custom-call.212 = (u32[2]{0}, u32[2]{0}) custom-call(u32[2]{0} %get-tuple-element.216, u32[2]{0} %get-tuple-element.217, u32[2]{0} %fusion.156, u32[2]{0} %fusion.155), custom_call_target="cuda_threefry2x32", operand_layout_constraints={u32[2]{0}, u32[2]{0}, u32[2]{0}, u32[2]{0}}, api_version=API_VERSION_STATUS_RETURNING, metadata={op_name="jit(training_step)/jit(main)/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}, backend_config="\002\000\000\000\000\000\000\000"
  %get-tuple-element.213 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}) %custom-call.212), index=0, metadata={op_name="jit(training_step)/jit(main)/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %get-tuple-element.214 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}) %custom-call.212), index=1, metadata={op_name="jit(training_step)/jit(main)/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %fusion.2 = (u32[2]{0}, u32[2]{0}, u32[2]{0}) fusion(u32[2]{0} %get-tuple-element.213, u32[2]{0} %get-tuple-element.214), kind=kLoop, calls=%fused_computation.2, metadata={op_name="jit(training_step)/jit(main)/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %get-tuple-element.215 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}, u32[2]{0}) %fusion.2), index=2
  %get-tuple-element.212 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}, u32[2]{0}) %fusion.2), index=1
  %custom-call.0 = (u32[2]{0}, u32[2]{0}) custom-call(u32[2]{0} %get-tuple-element.215, u32[2]{0} %get-tuple-element.212, u32[2]{0} %fusion.156, u32[2]{0} %fusion.155), custom_call_target="cuda_threefry2x32", operand_layout_constraints={u32[2]{0}, u32[2]{0}, u32[2]{0}, u32[2]{0}}, api_version=API_VERSION_STATUS_RETURNING, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}, backend_config="\002\000\000\000\000\000\000\000"
  %get-tuple-element.29 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}) %custom-call.0), index=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %get-tuple-element.30 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}) %custom-call.0), index=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %fusion.139 = (u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}) fusion(u32[2]{0} %get-tuple-element.29, u32[2]{0} %get-tuple-element.30), kind=kLoop, calls=%fused_computation.139, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %get-tuple-element.210 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}) %fusion.139), index=3
  %get-tuple-element.209 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}) %fusion.139), index=2
  %fusion.154 = u32[111]{0} fusion(), kind=kLoop, calls=%fused_computation.154, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(0,) limit_indices=(111,) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %fusion.143 = u32[111]{0} fusion(), kind=kLoop, calls=%fused_computation.143, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(111,) limit_indices=(222,) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %custom-call.2 = (u32[111]{0}, u32[111]{0}) custom-call(u32[111]{0} %get-tuple-element.210, u32[111]{0} %get-tuple-element.209, u32[111]{0} %fusion.154, u32[111]{0} %fusion.143), custom_call_target="cuda_threefry2x32", operand_layout_constraints={u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}}, api_version=API_VERSION_STATUS_RETURNING, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}, backend_config="o\000\000\000\000\000\000\000"
  %get-tuple-element.33 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}) %custom-call.2), index=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %get-tuple-element.34 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}) %custom-call.2), index=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %Arg_39.40 = f32[4096,3]{1,0} parameter(39)
  %Arg_38.39 = f32[4096,3]{1,0} parameter(38)
  %fusion.159 = (f32[4096,3]{1,0}, f32[4096,3]{1,0}) fusion(f32[4096,3]{1,0} %Arg_39.40, f32[4096,3]{1,0} %Arg_38.39, f32[2,3]{1,0} %Arg_34.35), kind=kLoop, calls=%fused_computation.159, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_max[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=368}
  %get-tuple-element.207 = f32[4096,3]{1,0} get-tuple-element((f32[4096,3]{1,0}, f32[4096,3]{1,0}) %fusion.159), index=0
  %get-tuple-element.208 = f32[4096,3]{1,0} get-tuple-element((f32[4096,3]{1,0}, f32[4096,3]{1,0}) %fusion.159), index=1
  %fusion.167 = (pred[4096]{0}, f32[4096]{0}, f32[4096]{0}) fusion(f32[4096,3]{1,0} %get-tuple-element.207, f32[4096,3]{1,0} %get-tuple-element.208), kind=kLoop, calls=%fused_computation.167, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=375}
  %get-tuple-element.204 = pred[4096]{0} get-tuple-element((pred[4096]{0}, f32[4096]{0}, f32[4096]{0}) %fusion.167), index=0
  %get-tuple-element.206 = f32[4096]{0} get-tuple-element((pred[4096]{0}, f32[4096]{0}, f32[4096]{0}) %fusion.167), index=2
  %get-tuple-element.205 = f32[4096]{0} get-tuple-element((pred[4096]{0}, f32[4096]{0}, f32[4096]{0}) %fusion.167), index=1
  %fusion.141 = f32[3,4096,221]{2,1,0} fusion(f32[2,3]{1,0} %Arg_34.35, u32[111]{0} %get-tuple-element.33, u32[111]{0} %get-tuple-element.34, f32[4096,3]{1,0} %Arg_39.40, f32[4096,3]{1,0} %Arg_38.39, /*index=5*/pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, f32[4096]{0} %get-tuple-element.205), kind=kLoop, calls=%fused_computation.141, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=128}
  %get-tuple-element.192 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}) %fusion.139), index=0
  %get-tuple-element.193 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}) %fusion.139), index=1
  %custom-call.1 = (u32[111]{0}, u32[111]{0}) custom-call(u32[111]{0} %get-tuple-element.192, u32[111]{0} %get-tuple-element.193, u32[111]{0} %fusion.154, u32[111]{0} %fusion.143), custom_call_target="cuda_threefry2x32", operand_layout_constraints={u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}}, api_version=API_VERSION_STATUS_RETURNING, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}, backend_config="o\000\000\000\000\000\000\000"
  %get-tuple-element.31 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}) %custom-call.1), index=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %get-tuple-element.32 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}) %custom-call.1), index=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %fusion.137 = f32[221]{0} fusion(u32[111]{0} %get-tuple-element.31, u32[111]{0} %get-tuple-element.32), kind=kLoop, calls=%fused_computation.137, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/log" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %Arg_9.10 = f32[3,16,128]{2,1,0} parameter(9)
  %fusion.130 = (f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}) fusion(f32[3,4096,221]{2,1,0} %fusion.141, s32[3,1]{1,0} %constant_367, s32[3,1]{1,0} %constant_363), kind=kLoop, calls=%fused_computation.130, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/squeeze[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %get-tuple-element.202 = f32[3,4096,221]{2,1,0} get-tuple-element((f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}) %fusion.130), index=1
  %fusion.131 = f32[3,16,4096,221]{3,2,1,0} fusion(f32[3,16,128]{2,1,0} %Arg_9.10, f32[3,4096,221]{2,1,0} %get-tuple-element.202), kind=kLoop, calls=%fused_computation.131, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(_map_coordinates))))/jit(jit_jvp(vmap(vmap(_map_coordinates))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %Arg_10.11 = f32[3,16,128,128]{3,2,1,0} parameter(10)
  %copy.53 = f32[3,16,128,128]{3,2,1,0} copy(f32[3,16,128,128]{3,2,1,0} %Arg_10.11)
  %get-tuple-element.203 = f32[3,4096,221]{2,1,0} get-tuple-element((f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}) %fusion.130), index=2
  %get-tuple-element.201 = f32[3,4096,221]{2,1,0} get-tuple-element((f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}) %fusion.130), index=0
  %input_fusion_reduce.6 = f32[4096,221]{1,0} fusion(f32[3,16,4096,221]{3,2,1,0} %fusion.131, f32[3,16,128,128]{3,2,1,0} %copy.53, f32[3,4096,221]{2,1,0} %get-tuple-element.203, f32[3,4096,221]{2,1,0} %get-tuple-element.201), kind=kInput, calls=%input_fused_computation_reduce.6, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %fusion.114 = f32[4096,221]{1,0} fusion(f32[4096,221]{1,0} %input_fusion_reduce.6), kind=kLoop, calls=%fused_computation.114, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %fusion.110 = f32[4096,55]{1,0} fusion(pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %fusion.114), kind=kLoop, calls=%fused_computation.110, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.108 = f32[4096,13]{1,0} fusion(f32[4096,55]{1,0} %fusion.110), kind=kLoop, calls=%fused_computation.108, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.106 = f32[4096,3]{1,0} fusion(f32[4096,13]{1,0} %fusion.108), kind=kLoop, calls=%fused_computation.106, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.105 = f32[4096,6]{1,0} fusion(f32[4096,3]{1,0} %fusion.106, f32[4096,13]{1,0} %fusion.108), kind=kLoop, calls=%fused_computation.105, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_70 = f32[] constant(0)
  %pad.59 = f32[4096,6]{1,0} pad(f32[4096,3]{1,0} %fusion.106, f32[] %constant_70), padding=0_0_0x1_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.161 = (f32[4096,27]{1,0}, f32[4096,27]{1,0}) fusion(f32[4096,13]{1,0} %fusion.108, f32[4096,6]{1,0} %fusion.105, f32[4096,6]{1,0} %pad.59, f32[4096,55]{1,0} %fusion.110), kind=kLoop, calls=%fused_computation.161, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %get-tuple-element.200 = f32[4096,27]{1,0} get-tuple-element((f32[4096,27]{1,0}, f32[4096,27]{1,0}) %fusion.161), index=1
  %get-tuple-element.199 = f32[4096,27]{1,0} get-tuple-element((f32[4096,27]{1,0}, f32[4096,27]{1,0}) %fusion.161), index=0
  %fusion.162 = (f32[4096,110]{1,0}, f32[4096,110]{1,0}) fusion(f32[4096,55]{1,0} %fusion.110, f32[4096,27]{1,0} %get-tuple-element.200, f32[4096,27]{1,0} %get-tuple-element.199, pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, /*index=5*/f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %fusion.114), kind=kLoop, calls=%fused_computation.162, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %get-tuple-element.198 = f32[4096,110]{1,0} get-tuple-element((f32[4096,110]{1,0}, f32[4096,110]{1,0}) %fusion.162), index=1
  %get-tuple-element.197 = f32[4096,110]{1,0} get-tuple-element((f32[4096,110]{1,0}, f32[4096,110]{1,0}) %fusion.162), index=0
  %fusion.98 = (f32[4096,221]{1,0}, f32[4096,221]{1,0}, f32[4096,221]{1,0}) fusion(f32[221]{0} %fusion.137, pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %fusion.114, /*index=5*/f32[4096,110]{1,0} %get-tuple-element.198, f32[4096,110]{1,0} %get-tuple-element.197), kind=kLoop, calls=%fused_computation.98, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %get-tuple-element.194 = f32[4096,221]{1,0} get-tuple-element((f32[4096,221]{1,0}, f32[4096,221]{1,0}, f32[4096,221]{1,0}) %fusion.98), index=0
  %iota.62 = s32[4096,221]{1,0} iota(), iota_dimension=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/broadcast_in_dim[shape=(4096, 221) broadcast_dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %sort.2 = (f32[4096,221]{1,0}, s32[4096,221]{1,0}, s32[4096,221]{1,0}) sort(f32[4096,221]{1,0} %get-tuple-element.194, s32[4096,221]{1,0} %iota.62, s32[4096,221]{1,0} %iota.62), dimensions={1}, is_stable=true, to_apply=%region_5.625.clone, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/sort[dimension=1 is_stable=True num_keys=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %get-tuple-element.51 = s32[4096,221]{1,0} get-tuple-element((f32[4096,221]{1,0}, s32[4096,221]{1,0}, s32[4096,221]{1,0}) %sort.2), index=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/sort[dimension=1 is_stable=True num_keys=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %fusion.97 = s32[4096,33,2]{1,0,2} fusion(s32[4096,221]{1,0} %get-tuple-element.51), kind=kLoop, calls=%fused_computation.97, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=2]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %fusion.96 = f32[3,3,4096,33]{3,2,0,1} fusion(s32[3,1]{1,0} %constant_367, s32[3,1]{1,0} %constant_363, f32[3,4096,221]{2,1,0} %fusion.141, s32[4096,33,2]{1,0,2} %fusion.97), kind=kLoop, calls=%fused_computation.96, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=162}
  %Arg_8.9 = f32[3,48,128,128]{3,2,1,0} parameter(8)
  %copy.51 = f32[3,48,128,128]{3,2,1,0} copy(f32[3,48,128,128]{3,2,1,0} %Arg_8.9)
  %fusion.75 = f32[135168,144]{0,1} fusion(f32[3,48,128]{2,1,0} %copy.50, f32[3,3,4096,33]{3,2,0,1} %fusion.96, f32[3,48,128,128]{3,2,1,0} %copy.51), kind=kLoop, calls=%fused_computation.75, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(135168, 144) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
  %cublas-gemm.1 = f32[135168,27]{1,0} custom-call(f32[135168,144]{0,1} %fusion.75, f32[144,27]{1,0} %copy.43), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=None]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"19464192\",\"rhs_stride\":\"3888\",\"selected_algorithm\":\"6\"}"
  %fusion.74 = f32[2]{0} fusion(), kind=kLoop, calls=%fused_computation.74, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/pow" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=26}
  %fusion.73 = f32[135168,27,4]{2,1,0} fusion(f32[135168,27]{1,0} %cublas-gemm.1, f32[2]{0} %fusion.74), kind=kLoop, calls=%fused_computation.73, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=2]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=31}
  %input_fusion_scatter.14 = u32[4096,221,3]{1,0,2} fusion(s32[4096,33,2]{1,0,2} %fusion.97), kind=kInput, calls=%input_fused_computation_scatter.14, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/scatter[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 1), scatter_dims_to_operand_dims=(0, 1)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %fusion.70 = pred[4096,33,3]{1,0,2} fusion(u32[4096,221,3]{1,0,2} %input_fusion_scatter.14, s32[4096,33,2]{1,0,2} %fusion.97), kind=kLoop, calls=%fused_computation.70, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %get-tuple-element.195 = f32[4096,221]{1,0} get-tuple-element((f32[4096,221]{1,0}, f32[4096,221]{1,0}, f32[4096,221]{1,0}) %fusion.98), index=1
  %fusion.72 = f32[135168,150]{1,0} fusion(f32[135168,27]{1,0} %cublas-gemm.1, f32[2]{0} %fusion.74, f32[135168,27,4]{2,1,0} %fusion.73, f32[4096,3]{1,0} %Arg_39.40), kind=kLoop, calls=%fused_computation.72, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=73}
  %Arg_2.3 = f32[150,128]{1,0} parameter(2)
  %copy.45 = f32[150,128]{1,0} copy(f32[150,128]{1,0} %Arg_2.3)
  %Arg_1.2 = f32[128]{0} parameter(1)
  %broadcast.545 = f32[135168,128]{1,0} broadcast(f32[128]{0} %Arg_1.2), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  %cublas-gemm.5 = f32[135168,128]{1,0} custom-call(f32[135168,150]{1,0} %fusion.72, f32[150,128]{1,0} %copy.45, f32[135168,128]{1,0} %broadcast.545), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":1,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"20275200\",\"rhs_stride\":\"19200\",\"selected_algorithm\":\"99\"}"
  %fusion.153 = f32[135168,128]{1,0} fusion(f32[135168,128]{1,0} %cublas-gemm.5), kind=kLoop, calls=%fused_computation.153, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(relu)/jit(jit_relu)/max" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  %Arg_4.5 = f32[128,128]{1,0} parameter(4)
  %copy.47 = f32[128,128]{1,0} copy(f32[128,128]{1,0} %Arg_4.5)
  %Arg_3.4 = f32[128]{0} parameter(3)
  %broadcast.548 = f32[135168,128]{1,0} broadcast(f32[128]{0} %Arg_3.4), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  %cublas-gemm.9 = f32[135168,128]{1,0} custom-call(f32[135168,128]{1,0} %fusion.153, f32[128,128]{1,0} %copy.47, f32[135168,128]{1,0} %broadcast.548), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":1,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"17301504\",\"rhs_stride\":\"16384\",\"selected_algorithm\":\"12\"}"
  %fusion.152 = f32[135168,128]{1,0} fusion(f32[135168,128]{1,0} %cublas-gemm.9), kind=kLoop, calls=%fused_computation.152, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(relu)/jit(jit_relu)/max" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=105}
  %Arg_6.7 = f32[128,3]{1,0} parameter(6)
  %copy.49 = f32[128,3]{1,0} copy(f32[128,3]{1,0} %Arg_6.7)
  %Arg_5.6 = f32[3]{0} parameter(5)
  %broadcast.551 = f32[135168,3]{0,1} broadcast(f32[3]{0} %Arg_5.6), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  %cublas-gemm.13 = f32[135168,3]{0,1} custom-call(f32[135168,128]{1,0} %fusion.152, f32[128,3]{1,0} %copy.49, f32[135168,3]{0,1} %broadcast.551), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":1,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"17301504\",\"rhs_stride\":\"384\",\"selected_algorithm\":\"11\"}"
  %Arg_37.38 = f32[4096,3]{1,0} parameter(37)
  %input_fusion_scatter.13 = f32[4096,221,3]{1,0,2} fusion(s32[4096,33,2]{1,0,2} %fusion.97, pred[4096,33,3]{1,0,2} %fusion.70, f32[135168,3]{0,1} %cublas-gemm.13), kind=kInput, calls=%input_fused_computation_scatter.13, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 1), scatter_dims_to_operand_dims=(0, 1)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %input_fusion_reduce.5 = f32[3,4096]{0,1} fusion(f32[4096,221,3]{1,0,2} %input_fusion_scatter.13, f32[4096,221]{1,0} %get-tuple-element.195), kind=kInput, calls=%input_fused_computation_reduce.5
  %get-tuple-element.196 = f32[4096,221]{1,0} get-tuple-element((f32[4096,221]{1,0}, f32[4096,221]{1,0}, f32[4096,221]{1,0}) %fusion.98), index=2
  %input_fusion_reduce.4 = f32[4096]{0} fusion(f32[4096,221]{1,0} %get-tuple-element.195, s32[4096,33,2]{1,0,2} %fusion.97), kind=kInput, calls=%input_fused_computation_reduce.4, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  %fusion.66 = f32[4096,3]{1,0} fusion(f32[4096,3]{1,0} %Arg_37.38, f32[3,4096]{0,1} %input_fusion_reduce.5, f32[4096,221]{1,0} %get-tuple-element.196, f32[4096]{0} %input_fusion_reduce.4), kind=kLoop, calls=%fused_computation.66, metadata={op_name="jit(training_step)/jit(main)/sub" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %fusion.64 = f32[135168,3]{0,1} fusion(pred[4096,33,3]{1,0,2} %fusion.70, s32[4096,33,2]{1,0,2} %fusion.97, f32[4096,221]{1,0} %get-tuple-element.195, f32[135168,3]{0,1} %cublas-gemm.13, f32[4096,3]{1,0} %fusion.66, /*index=5*/f32[4096]{0} %input_fusion_reduce.4, f32[4096,221]{1,0} %get-tuple-element.196), kind=kLoop, calls=%fused_computation.64, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(sigmoid)))/jit(jit_transpose(jvp(sigmoid)))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %cublas-gemm.15 = f32[135168,128]{1,0} custom-call(f32[135168,3]{0,1} %fusion.64, f32[128,3]{1,0} %copy.49), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/dot_general[dimension_numbers=(((1,), (1,)), ((), ())) precision=None preferred_element_type=None]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"1\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"405504\",\"rhs_stride\":\"384\",\"selected_algorithm\":\"-1\"}"
  %fusion.63 = f32[135168,128]{1,0} fusion(f32[135168,128]{1,0} %cublas-gemm.15, f32[135168,128]{1,0} %cublas-gemm.9), kind=kLoop, calls=%fused_computation.63, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=105}
  %cublas-gemm.17 = f32[135168,128]{1,0} custom-call(f32[135168,128]{1,0} %fusion.63, f32[128,128]{1,0} %copy.47), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/dot_general[dimension_numbers=(((1,), (1,)), ((), ())) precision=None preferred_element_type=None]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"1\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"17301504\",\"rhs_stride\":\"16384\",\"selected_algorithm\":\"-1\"}"
  %fusion.62 = f32[135168,128]{1,0} fusion(f32[135168,128]{1,0} %cublas-gemm.17, f32[135168,128]{1,0} %cublas-gemm.5), kind=kLoop, calls=%fused_computation.62, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  %cublas-gemm.19 = f32[135168,150]{1,0} custom-call(f32[135168,128]{1,0} %fusion.62, f32[150,128]{1,0} %copy.45), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/dot_general[dimension_numbers=(((1,), (1,)), ((), ())) precision=None preferred_element_type=None]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"1\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"17301504\",\"rhs_stride\":\"19200\",\"selected_algorithm\":\"-1\"}"
  %fusion.61 = f32[135168,27]{1,0} fusion(f32[135168,27,4]{2,1,0} %fusion.73, f32[135168,150]{1,0} %cublas-gemm.19, f32[2]{0} %fusion.74), kind=kLoop, calls=%fused_computation.61, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %cublas-gemm.21 = f32[144,27]{1,0} custom-call(f32[135168,144]{0,1} %fusion.75, f32[135168,27]{1,0} %fusion.61), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/transpose[permutation=(1, 0)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"0\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"19464192\",\"rhs_stride\":\"3649536\",\"selected_algorithm\":\"0\"}"
  %Arg_12.13 = f32[144,27]{1,0} parameter(12)
  %fusion.54 = (f32[144,27]{1,0}, f32[144,27]{1,0}, f32[144,27]{1,0}) fusion(f32[144,27]{1,0} %copy.43, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[144,27]{1,0} %Arg_23.24, /*index=5*/f32[144,27]{1,0} %cublas-gemm.21, f32[144,27]{1,0} %Arg_12.13), kind=kLoop, calls=%fused_computation.54, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.222 = f32[144,27]{1,0} get-tuple-element((f32[144,27]{1,0}, f32[144,27]{1,0}, f32[144,27]{1,0}) %fusion.54), index=0
  %Arg_26.27 = f32[128]{0} parameter(26)
  %input_fusion_reduce.2 = (f32[368,128]{1,0}, f32[368,128]{1,0}) fusion(f32[135168,128]{1,0} %fusion.63, f32[135168,128]{1,0} %fusion.62), kind=kInput, calls=%input_fused_computation_reduce.2
  %get-tuple-element.220 = f32[368,128]{1,0} get-tuple-element((f32[368,128]{1,0}, f32[368,128]{1,0}) %input_fusion_reduce.2), index=0
  %get-tuple-element.221 = f32[368,128]{1,0} get-tuple-element((f32[368,128]{1,0}, f32[368,128]{1,0}) %input_fusion_reduce.2), index=1
  %input_fusion_reduce.7 = (f32[128]{0}, f32[128]{0}) fusion(f32[368,128]{1,0} %get-tuple-element.220, f32[] %constant_70, f32[368,128]{1,0} %get-tuple-element.221), kind=kInput, calls=%input_fused_computation_reduce.7, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  %get-tuple-element.218 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}) %input_fusion_reduce.7), index=0
  %Arg_15.16 = f32[128]{0} parameter(15)
  %Arg_24.25 = f32[128]{0} parameter(24)
  %get-tuple-element.219 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}) %input_fusion_reduce.7), index=1
  %Arg_13.14 = f32[128]{0} parameter(13)
  %fusion.45 = (f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) fusion(f32[128]{0} %Arg_3.4, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[128]{0} %Arg_26.27, /*index=5*/f32[128]{0} %get-tuple-element.218, f32[128]{0} %Arg_15.16, f32[128]{0} %Arg_1.2, f32[128]{0} %Arg_24.25, f32[128]{0} %get-tuple-element.219, /*index=10*/f32[128]{0} %Arg_13.14), kind=kLoop, calls=%fused_computation.45, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.223 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) %fusion.45), index=3
  %Arg_25.26 = f32[150,128]{1,0} parameter(25)
  %cublas-gemm.23 = f32[150,128]{1,0} custom-call(f32[135168,150]{1,0} %fusion.72, f32[135168,128]{1,0} %fusion.62), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/transpose[permutation=(1, 0)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"0\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"20275200\",\"rhs_stride\":\"17301504\",\"selected_algorithm\":\"0\"}"
  %Arg_14.15 = f32[150,128]{1,0} parameter(14)
  %fusion.48 = (f32[150,128]{1,0}, f32[150,128]{1,0}, f32[150,128]{1,0}) fusion(f32[150,128]{1,0} %copy.45, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[150,128]{1,0} %Arg_25.26, /*index=5*/f32[150,128]{1,0} %cublas-gemm.23, f32[150,128]{1,0} %Arg_14.15), kind=kLoop, calls=%fused_computation.48, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.224 = f32[150,128]{1,0} get-tuple-element((f32[150,128]{1,0}, f32[150,128]{1,0}, f32[150,128]{1,0}) %fusion.48), index=0
  %get-tuple-element.225 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) %fusion.45), index=0
  %Arg_27.28 = f32[128,128]{1,0} parameter(27)
  %cublas-gemm.25 = f32[128,128]{1,0} custom-call(f32[135168,128]{1,0} %fusion.153, f32[135168,128]{1,0} %fusion.63), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/transpose[permutation=(1, 0)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"0\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"17301504\",\"rhs_stride\":\"17301504\",\"selected_algorithm\":\"11\"}"
  %Arg_16.17 = f32[128,128]{1,0} parameter(16)
  %fusion.42 = (f32[128,128]{1,0}, f32[128,128]{1,0}, f32[128,128]{1,0}) fusion(f32[128,128]{1,0} %copy.47, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[128,128]{1,0} %Arg_27.28, /*index=5*/f32[128,128]{1,0} %cublas-gemm.25, f32[128,128]{1,0} %Arg_16.17), kind=kLoop, calls=%fused_computation.42, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.226 = f32[128,128]{1,0} get-tuple-element((f32[128,128]{1,0}, f32[128,128]{1,0}, f32[128,128]{1,0}) %fusion.42), index=0
  %Arg_28.29 = f32[3]{0} parameter(28)
  %input_fusion_reduce.1 = f32[3,368]{1,0} fusion(f32[135168,3]{0,1} %fusion.64), kind=kInput, calls=%input_fused_computation_reduce.1
  %reduce.17 = f32[3]{0} reduce(f32[3,368]{1,0} %input_fusion_reduce.1, f32[] %constant_70), dimensions={1}, to_apply=%region_27.1710, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  %Arg_17.18 = f32[3]{0} parameter(17)
  %fusion.39 = (f32[3]{0}, f32[3]{0}, f32[3]{0}) fusion(f32[3]{0} %Arg_5.6, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[3]{0} %Arg_28.29, /*index=5*/f32[3]{0} %reduce.17, f32[3]{0} %Arg_17.18), kind=kLoop, calls=%fused_computation.39, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.227 = f32[3]{0} get-tuple-element((f32[3]{0}, f32[3]{0}, f32[3]{0}) %fusion.39), index=0
  %Arg_29.30 = f32[128,3]{1,0} parameter(29)
  %cublas-gemm.27 = f32[128,3]{1,0} custom-call(f32[135168,128]{1,0} %fusion.152, f32[135168,3]{0,1} %fusion.64), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/transpose[permutation=(1, 0)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"0\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"17301504\",\"rhs_stride\":\"405504\",\"selected_algorithm\":\"-1\"}"
  %Arg_18.19 = f32[128,3]{1,0} parameter(18)
  %fusion.36 = (f32[128,3]{1,0}, f32[128,3]{1,0}, f32[128,3]{1,0}) fusion(f32[128,3]{1,0} %copy.49, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[128,3]{1,0} %Arg_29.30, /*index=5*/f32[128,3]{1,0} %cublas-gemm.27, f32[128,3]{1,0} %Arg_18.19), kind=kLoop, calls=%fused_computation.36, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.228 = f32[128,3]{1,0} get-tuple-element((f32[128,3]{1,0}, f32[128,3]{1,0}, f32[128,3]{1,0}) %fusion.36), index=0
  %Arg_30.31 = f32[3,48,128]{2,1,0} parameter(30)
  %cublas-gemm.29 = f32[135168,144]{1,0} custom-call(f32[135168,27]{1,0} %fusion.61, f32[144,27]{1,0} %copy.43), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/dot_general[dimension_numbers=(((1,), (1,)), ((), ())) precision=None preferred_element_type=None]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"1\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"3649536\",\"rhs_stride\":\"3888\",\"selected_algorithm\":\"0\"}"
  %fusion.34 = f32[3,48,4096,33]{2,1,0,3} fusion(f32[3,48,128,128]{3,2,1,0} %copy.51, f32[3,3,4096,33]{3,2,0,1} %fusion.96, f32[135168,144]{1,0} %cublas-gemm.29), kind=kLoop, calls=%fused_computation.34, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %input_fusion_scatter.12 = f32[3,48,128]{2,1,0} fusion(f32[3,48,4096,33]{2,1,0,3} %fusion.34, f32[3,3,4096,33]{3,2,0,1} %fusion.96), kind=kInput, calls=%input_fused_computation_scatter.12, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.11 = f32[3,48,128]{2,1,0} fusion(f32[3,48,4096,33]{2,1,0,3} %fusion.34, f32[3,3,4096,33]{3,2,0,1} %fusion.96), kind=kInput, calls=%input_fused_computation_scatter.11, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %Arg_19.20 = f32[3,48,128]{2,1,0} parameter(19)
  %fusion.31 = (f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}) fusion(f32[3,48,128]{2,1,0} %copy.50, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[3,48,128]{2,1,0} %Arg_30.31, /*index=5*/f32[3,48,128]{2,1,0} %input_fusion_scatter.12, f32[3,48,128]{2,1,0} %input_fusion_scatter.11, f32[3,48,128]{2,1,0} %Arg_19.20), kind=kLoop, calls=%fused_computation.31, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.229 = f32[3,48,128]{2,1,0} get-tuple-element((f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}) %fusion.31), index=0
  %Arg_31.32 = f32[3,48,128,128]{3,2,1,0} parameter(31)
  %fusion.30 = f32[3,48,4096,33]{3,2,1,0} fusion(f32[3,48,128]{2,1,0} %copy.50, f32[3,3,4096,33]{3,2,0,1} %fusion.96, f32[135168,144]{1,0} %cublas-gemm.29), kind=kLoop, calls=%fused_computation.30, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %input_fusion_scatter.10 = f32[3,48,128,128]{3,2,1,0} fusion(f32[3,48,4096,33]{3,2,1,0} %fusion.30, f32[3,3,4096,33]{3,2,0,1} %fusion.96), kind=kInput, calls=%input_fused_computation_scatter.10, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.9 = f32[3,48,128,128]{3,2,1,0} fusion(f32[3,48,4096,33]{3,2,1,0} %fusion.30, f32[3,3,4096,33]{3,2,0,1} %fusion.96), kind=kInput, calls=%input_fused_computation_scatter.9, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.8 = f32[3,48,128,128]{3,2,1,0} fusion(f32[3,48,4096,33]{3,2,1,0} %fusion.30, f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.10, f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.9, f32[3,3,4096,33]{3,2,0,1} %fusion.96), kind=kInput, calls=%input_fused_computation_scatter.8, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.7 = f32[3,48,128,128]{3,2,1,0} fusion(f32[3,48,4096,33]{3,2,1,0} %fusion.30, f32[3,3,4096,33]{3,2,0,1} %fusion.96), kind=kInput, calls=%input_fused_computation_scatter.7, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %Arg_20.21 = f32[3,48,128,128]{3,2,1,0} parameter(20)
  %fusion.27 = (f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}) fusion(f32[3,48,128,128]{3,2,1,0} %copy.51, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[3,48,128,128]{3,2,1,0} %Arg_31.32, /*index=5*/f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.8, f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.7, f32[3,48,128,128]{3,2,1,0} %Arg_20.21), kind=kLoop, calls=%fused_computation.27, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.230 = f32[3,48,128,128]{3,2,1,0} get-tuple-element((f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}) %fusion.27), index=0
  %Arg_32.33 = f32[3,16,128]{2,1,0} parameter(32)
  %fusion.24 = f32[4096,221]{1,0} fusion(f32[4096,221,3]{1,0,2} %input_fusion_scatter.13, f32[4096,3]{1,0} %fusion.66, f32[4096]{0} %input_fusion_reduce.4, f32[4096,221]{1,0} %get-tuple-element.196), kind=kLoop, calls=%fused_computation.24, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %fusion.25 = f32[4096]{0} fusion(f32[3,4096]{0,1} %input_fusion_reduce.5, f32[4096,3]{1,0} %fusion.66), kind=kLoop, calls=%fused_computation.25, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %input_fusion_scatter.6 = f32[4096,221]{1,0} fusion(f32[4096,221]{1,0} %fusion.24, s32[4096,33,2]{1,0,2} %fusion.97, f32[4096]{0} %fusion.25, f32[4096,221]{1,0} %get-tuple-element.196, f32[4096]{0} %input_fusion_reduce.4), kind=kInput, calls=%input_fused_computation_scatter.6, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(), inserted_window_dims=(0, 1), scatter_dims_to_operand_dims=(0, 1)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=472}
  %fusion.23 = f32[4096,221]{1,0} fusion(f32[4096,221]{1,0} %get-tuple-element.196, f32[4096,221]{1,0} %input_fusion_scatter.6, f32[4096]{0} %fusion.25, pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, /*index=5*/f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %fusion.114, f32[4096]{0} %input_fusion_reduce.4, f32[4096,3]{1,0} %fusion.66), kind=kLoop, calls=%fused_computation.23, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.21 = f32[4096,55]{1,0} fusion(f32[4096,221]{1,0} %fusion.23), kind=kLoop, calls=%fused_computation.21, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.19 = f32[4096,13]{1,0} fusion(f32[4096,55]{1,0} %fusion.21), kind=kLoop, calls=%fused_computation.19, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.17 = f32[4096,3]{1,0} fusion(f32[4096,13]{1,0} %fusion.19), kind=kLoop, calls=%fused_computation.17, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.16 = f32[4096,6]{1,0} fusion(f32[4096,3]{1,0} %fusion.17, f32[4096,13]{1,0} %fusion.19), kind=kLoop, calls=%fused_computation.16, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.73 = f32[4096,6]{1,0} pad(f32[4096,3]{1,0} %fusion.17, f32[] %constant_70), padding=0_0_0x1_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.164 = (f32[4096,27]{1,0}, f32[4096,27]{1,0}) fusion(f32[4096,13]{1,0} %fusion.19, f32[4096,6]{1,0} %fusion.16, f32[4096,6]{1,0} %pad.73, f32[4096,55]{1,0} %fusion.21), kind=kLoop, calls=%fused_computation.164, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %get-tuple-element.155 = f32[4096,27]{1,0} get-tuple-element((f32[4096,27]{1,0}, f32[4096,27]{1,0}) %fusion.164), index=1
  %get-tuple-element.154 = f32[4096,27]{1,0} get-tuple-element((f32[4096,27]{1,0}, f32[4096,27]{1,0}) %fusion.164), index=0
  %fusion.165 = (f32[4096,110]{1,0}, f32[4096,110]{1,0}) fusion(f32[4096,55]{1,0} %fusion.21, f32[4096,27]{1,0} %get-tuple-element.155, f32[4096,27]{1,0} %get-tuple-element.154, f32[4096,221]{1,0} %fusion.23), kind=kLoop, calls=%fused_computation.165, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %get-tuple-element.153 = f32[4096,110]{1,0} get-tuple-element((f32[4096,110]{1,0}, f32[4096,110]{1,0}) %fusion.165), index=1
  %get-tuple-element.152 = f32[4096,110]{1,0} get-tuple-element((f32[4096,110]{1,0}, f32[4096,110]{1,0}) %fusion.165), index=0
  %fusion.10 = f32[3,16,4096,221]{2,1,0,3} fusion(f32[3,16,128,128]{3,2,1,0} %copy.53, f32[3,4096,221]{2,1,0} %get-tuple-element.203, f32[3,4096,221]{2,1,0} %get-tuple-element.201, f32[4096,221]{1,0} %fusion.114, f32[4096,221]{1,0} %input_fusion_reduce.6, /*index=5*/pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %input_fusion_scatter.6, f32[4096,221]{1,0} %get-tuple-element.196, /*index=10*/f32[4096,221]{1,0} %fusion.23, f32[4096,110]{1,0} %get-tuple-element.153, f32[4096,110]{1,0} %get-tuple-element.152), kind=kLoop, calls=%fused_computation.10, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %input_fusion_scatter.5 = f32[3,16,128]{2,1,0} fusion(f32[3,16,4096,221]{2,1,0,3} %fusion.10, f32[3,4096,221]{2,1,0} %get-tuple-element.202), kind=kInput, calls=%input_fused_computation_scatter.5, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.4 = f32[3,16,128]{2,1,0} fusion(f32[3,16,4096,221]{2,1,0,3} %fusion.10, f32[3,4096,221]{2,1,0} %get-tuple-element.202), kind=kInput, calls=%input_fused_computation_scatter.4, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %Arg_21.22 = f32[3,16,128]{2,1,0} parameter(21)
  %fusion.7 = (f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}) fusion(f32[3,16,128]{2,1,0} %Arg_9.10, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[3,16,128]{2,1,0} %Arg_32.33, /*index=5*/f32[3,16,128]{2,1,0} %input_fusion_scatter.5, f32[3,16,128]{2,1,0} %input_fusion_scatter.4, f32[3,16,128]{2,1,0} %Arg_21.22), kind=kLoop, calls=%fused_computation.7, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.231 = f32[3,16,128]{2,1,0} get-tuple-element((f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}) %fusion.7), index=0
  %Arg_33.34 = f32[3,16,128,128]{3,2,1,0} parameter(33)
  %input_fusion_scatter.3 = f32[3,16,128,128]{3,2,1,0} fusion(f32[3,4096,221]{2,1,0} %get-tuple-element.203, f32[3,4096,221]{2,1,0} %get-tuple-element.201, f32[3,16,4096,221]{3,2,1,0} %fusion.131, f32[4096,221]{1,0} %fusion.114, f32[4096,221]{1,0} %input_fusion_reduce.6, /*index=5*/pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %input_fusion_scatter.6, f32[4096,221]{1,0} %get-tuple-element.196, /*index=10*/f32[4096,221]{1,0} %fusion.23, f32[4096,110]{1,0} %get-tuple-element.153, f32[4096,110]{1,0} %get-tuple-element.152), kind=kInput, calls=%input_fused_computation_scatter.3, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.2 = f32[3,16,128,128]{3,2,1,0} fusion(f32[3,4096,221]{2,1,0} %get-tuple-element.203, f32[3,4096,221]{2,1,0} %get-tuple-element.201, f32[3,16,4096,221]{3,2,1,0} %fusion.131, f32[4096,221]{1,0} %fusion.114, f32[4096,221]{1,0} %input_fusion_reduce.6, /*index=5*/pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %input_fusion_scatter.6, f32[4096,221]{1,0} %get-tuple-element.196, /*index=10*/f32[4096,221]{1,0} %fusion.23, f32[4096,110]{1,0} %get-tuple-element.153, f32[4096,110]{1,0} %get-tuple-element.152), kind=kInput, calls=%input_fused_computation_scatter.2, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.1 = f32[3,16,128,128]{3,2,1,0} fusion(f32[3,16,128,128]{3,2,1,0} %input_fusion_scatter.3, f32[3,16,128,128]{3,2,1,0} %input_fusion_scatter.2, f32[3,4096,221]{2,1,0} %get-tuple-element.203, f32[3,4096,221]{2,1,0} %get-tuple-element.201, f32[3,16,4096,221]{3,2,1,0} %fusion.131, /*index=5*/f32[4096,221]{1,0} %fusion.114, f32[4096,221]{1,0} %input_fusion_reduce.6, pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, f32[4096]{0} %get-tuple-element.205, /*index=10*/f32[4096,221]{1,0} %input_fusion_scatter.6, f32[4096,221]{1,0} %get-tuple-element.196, f32[4096,221]{1,0} %fusion.23, f32[4096,110]{1,0} %get-tuple-element.153, f32[4096,110]{1,0} %get-tuple-element.152), kind=kInput, calls=%input_fused_computation_scatter.1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter = f32[3,16,128,128]{3,2,1,0} fusion(f32[3,4096,221]{2,1,0} %get-tuple-element.203, f32[3,4096,221]{2,1,0} %get-tuple-element.201, f32[3,16,4096,221]{3,2,1,0} %fusion.131, f32[4096,221]{1,0} %fusion.114, f32[4096,221]{1,0} %input_fusion_reduce.6, /*index=5*/pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %input_fusion_scatter.6, f32[4096,221]{1,0} %get-tuple-element.196, /*index=10*/f32[4096,221]{1,0} %fusion.23, f32[4096,110]{1,0} %get-tuple-element.153, f32[4096,110]{1,0} %get-tuple-element.152), kind=kInput, calls=%input_fused_computation_scatter, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(_map_coordinates)))))/jit(jit_transpose(jvp(vmap(vmap(_map_coordinates)))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(1,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %Arg_22.23 = f32[3,16,128,128]{3,2,1,0} parameter(22)
  %fusion.3 = (f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}) fusion(f32[3,16,128,128]{3,2,1,0} %copy.53, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[3,16,128,128]{3,2,1,0} %Arg_33.34, /*index=5*/f32[3,16,128,128]{3,2,1,0} %input_fusion_scatter.1, f32[3,16,128,128]{3,2,1,0} %input_fusion_scatter, f32[3,16,128,128]{3,2,1,0} %Arg_22.23), kind=kLoop, calls=%fused_computation.3, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.232 = f32[3,16,128,128]{3,2,1,0} get-tuple-element((f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}) %fusion.3), index=0
  %get-tuple-element.233 = s32[] get-tuple-element((f32[], f32[], s32[]) %fusion.56), index=2
  %get-tuple-element.234 = f32[144,27]{1,0} get-tuple-element((f32[144,27]{1,0}, f32[144,27]{1,0}, f32[144,27]{1,0}) %fusion.54), index=2
  %get-tuple-element.235 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) %fusion.45), index=5
  %get-tuple-element.236 = f32[150,128]{1,0} get-tuple-element((f32[150,128]{1,0}, f32[150,128]{1,0}, f32[150,128]{1,0}) %fusion.48), index=2
  %get-tuple-element.237 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) %fusion.45), index=2
  %get-tuple-element.238 = f32[128,128]{1,0} get-tuple-element((f32[128,128]{1,0}, f32[128,128]{1,0}, f32[128,128]{1,0}) %fusion.42), index=2
  %get-tuple-element.239 = f32[3]{0} get-tuple-element((f32[3]{0}, f32[3]{0}, f32[3]{0}) %fusion.39), index=2
  %get-tuple-element.240 = f32[128,3]{1,0} get-tuple-element((f32[128,3]{1,0}, f32[128,3]{1,0}, f32[128,3]{1,0}) %fusion.36), index=2
  %get-tuple-element.241 = f32[3,48,128]{2,1,0} get-tuple-element((f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}) %fusion.31), index=2
  %get-tuple-element.242 = f32[3,48,128,128]{3,2,1,0} get-tuple-element((f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}) %fusion.27), index=2
  %get-tuple-element.243 = f32[3,16,128]{2,1,0} get-tuple-element((f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}) %fusion.7), index=2
  %get-tuple-element.244 = f32[3,16,128,128]{3,2,1,0} get-tuple-element((f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}) %fusion.3), index=2
  %get-tuple-element.245 = f32[144,27]{1,0} get-tuple-element((f32[144,27]{1,0}, f32[144,27]{1,0}, f32[144,27]{1,0}) %fusion.54), index=1
  %get-tuple-element.246 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) %fusion.45), index=4
  %get-tuple-element.247 = f32[150,128]{1,0} get-tuple-element((f32[150,128]{1,0}, f32[150,128]{1,0}, f32[150,128]{1,0}) %fusion.48), index=1
  %get-tuple-element.248 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) %fusion.45), index=1
  %get-tuple-element.249 = f32[128,128]{1,0} get-tuple-element((f32[128,128]{1,0}, f32[128,128]{1,0}, f32[128,128]{1,0}) %fusion.42), index=1
  %get-tuple-element.250 = f32[3]{0} get-tuple-element((f32[3]{0}, f32[3]{0}, f32[3]{0}) %fusion.39), index=1
  %get-tuple-element.251 = f32[128,3]{1,0} get-tuple-element((f32[128,3]{1,0}, f32[128,3]{1,0}, f32[128,3]{1,0}) %fusion.36), index=1
  %get-tuple-element.252 = f32[3,48,128]{2,1,0} get-tuple-element((f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}) %fusion.31), index=1
  %get-tuple-element.253 = f32[3,48,128,128]{3,2,1,0} get-tuple-element((f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}) %fusion.27), index=1
  %get-tuple-element.254 = f32[3,16,128]{2,1,0} get-tuple-element((f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}) %fusion.7), index=1
  %get-tuple-element.255 = f32[3,16,128,128]{3,2,1,0} get-tuple-element((f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}) %fusion.3), index=1
  %get-tuple-element.257 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}, u32[2]{0}) %fusion.2), index=0
  %get-tuple-element.258 = s32[] get-tuple-element((f32[], f32[], s32[]) %fusion.166), index=2
  %get-tuple-element.259 = f32[] get-tuple-element((f32[], f32[], s32[]) %fusion.166), index=0
  %constant_43 = f32[] constant(0.02)
  %multiply.2257 = f32[] multiply(f32[] %get-tuple-element.184, f32[] %constant_43), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=355}
  %input_fusion_reduce = f32[] fusion(f32[4096,3]{1,0} %fusion.66), kind=kInput, calls=%input_fused_computation_reduce, metadata={op_name="jit(training_step)/jit(main)/reduce_sum[axes=(0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %fusion.1 = (f32[], f32[]) fusion(f32[] %input_fusion_reduce), kind=kLoop, calls=%fused_computation.1, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/utils.py" source_line=17}
  %get-tuple-element.261 = f32[] get-tuple-element((f32[], f32[]) %fusion.1), index=1
  %get-tuple-element.262 = f32[] get-tuple-element((f32[], f32[]) %fusion.1), index=0
  ROOT %tuple.91 = (f32[144,27]{1,0}, f32[128]{0}, f32[150,128]{1,0}, f32[128]{0}, f32[128,128]{1,0}, /*index=5*/f32[3]{0}, f32[128,3]{1,0}, f32[3,48,128]{2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,16,128]{2,1,0}, /*index=10*/f32[3,16,128,128]{3,2,1,0}, s32[], f32[144,27]{1,0}, f32[128]{0}, f32[150,128]{1,0}, /*index=15*/f32[128]{0}, f32[128,128]{1,0}, f32[3]{0}, f32[128,3]{1,0}, f32[3,48,128]{2,1,0}, /*index=20*/f32[3,48,128,128]{3,2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[144,27]{1,0}, f32[128]{0}, /*index=25*/f32[150,128]{1,0}, f32[128]{0}, f32[128,128]{1,0}, f32[3]{0}, f32[128,3]{1,0}, /*index=30*/f32[3,48,128]{2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[2,3]{1,0}, /*index=35*/u32[2]{0}, s32[], f32[], f32[], f32[], /*index=40*/f32[]) tuple(f32[144,27]{1,0} %get-tuple-element.222, f32[128]{0} %get-tuple-element.223, f32[150,128]{1,0} %get-tuple-element.224, f32[128]{0} %get-tuple-element.225, f32[128,128]{1,0} %get-tuple-element.226, /*index=5*/f32[3]{0} %get-tuple-element.227, f32[128,3]{1,0} %get-tuple-element.228, f32[3,48,128]{2,1,0} %get-tuple-element.229, f32[3,48,128,128]{3,2,1,0} %get-tuple-element.230, f32[3,16,128]{2,1,0} %get-tuple-element.231, /*index=10*/f32[3,16,128,128]{3,2,1,0} %get-tuple-element.232, s32[] %get-tuple-element.233, f32[144,27]{1,0} %get-tuple-element.234, f32[128]{0} %get-tuple-element.235, f32[150,128]{1,0} %get-tuple-element.236, /*index=15*/f32[128]{0} %get-tuple-element.237, f32[128,128]{1,0} %get-tuple-element.238, f32[3]{0} %get-tuple-element.239, f32[128,3]{1,0} %get-tuple-element.240, f32[3,48,128]{2,1,0} %get-tuple-element.241, /*index=20*/f32[3,48,128,128]{3,2,1,0} %get-tuple-element.242, f32[3,16,128]{2,1,0} %get-tuple-element.243, f32[3,16,128,128]{3,2,1,0} %get-tuple-element.244, f32[144,27]{1,0} %get-tuple-element.245, f32[128]{0} %get-tuple-element.246, /*index=25*/f32[150,128]{1,0} %get-tuple-element.247, f32[128]{0} %get-tuple-element.248, f32[128,128]{1,0} %get-tuple-element.249, f32[3]{0} %get-tuple-element.250, f32[128,3]{1,0} %get-tuple-element.251, /*index=30*/f32[3,48,128]{2,1,0} %get-tuple-element.252, f32[3,48,128,128]{3,2,1,0} %get-tuple-element.253, f32[3,16,128]{2,1,0} %get-tuple-element.254, f32[3,16,128,128]{3,2,1,0} %get-tuple-element.255, f32[2,3]{1,0} %Arg_34.35, /*index=35*/u32[2]{0} %get-tuple-element.257, s32[] %get-tuple-element.258, f32[] %get-tuple-element.259, f32[] %multiply.2257, f32[] %get-tuple-element.261, /*index=40*/f32[] %get-tuple-element.262)
}

Texto alterado

Abrir arquivo

%region_5.628.clone (Arg_0.0: f32[], Arg_1.0: f32[], Arg_2.0: s32[], Arg_3.0: s32[], p.2.lhs.clone: s32[], p.2.rhs.clone: s32[]) -> pred[] {
  %Arg_2.0 = s32[] parameter(2)
  %Arg_3.0 = s32[] parameter(3)
  %p.2.lhs.clone = s32[] parameter(4)
  %p.2.rhs.clone = s32[] parameter(5)
  %Arg_1.0 = f32[] parameter(1)
  %Arg_0.0 = f32[] parameter(0)
  ROOT %fusion = pred[] fusion(s32[] %p.2.lhs.clone, s32[] %p.2.rhs.clone, f32[] %Arg_1.0, f32[] %Arg_0.0), kind=kLoop, calls=%fused_computation
}

%region_27.1717 (Arg_0.1718: f32[], Arg_1.1719: f32[]) -> f32[] {
  %Arg_0.1718 = f32[] parameter(0)
  %Arg_1.1719 = f32[] parameter(1)
  ROOT %add.1720 = f32[] add(f32[] %Arg_0.1718, f32[] %Arg_1.1719), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
}

%fused_computation.1 (param_0.369: f32[]) -> (f32[], f32[]) {
  %param_0.369 = f32[] parameter(0)
  %constant_471 = f32[] constant(8.13802108e-05)
  %multiply.396 = f32[] multiply(f32[] %param_0.369, f32[] %constant_471), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %constant_454 = f32[] constant(1e-10)
  %maximum.59 = f32[] maximum(f32[] %multiply.396, f32[] %constant_454), metadata={op_name="jit(training_step)/jit(main)/max" source_file="/home/brent/tensorf-jax/tensorf/utils.py" source_line=8}
  %log.3 = f32[] log(f32[] %maximum.59), metadata={op_name="jit(training_step)/jit(main)/log" source_file="/home/brent/tensorf-jax/tensorf/utils.py" source_line=17}
  %constant_453 = f32[] constant(-4.34294462)
  %multiply.232 = f32[] multiply(f32[] %log.3, f32[] %constant_453), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/utils.py" source_line=17}
  ROOT %tuple.44 = (f32[], f32[]) tuple(f32[] %multiply.232, f32[] %multiply.396)
}

%region_10.1472 (Arg_0.1473: f32[], Arg_1.1474: f32[]) -> f32[] {
  %Arg_0.1473 = f32[] parameter(0)
  %Arg_1.1474 = f32[] parameter(1)
  ROOT %add.1475 = f32[] add(f32[] %Arg_0.1473, f32[] %Arg_1.1474), metadata={op_name="jit(training_step)/jit(main)/reduce_sum[axes=(0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
}

%input_fused_computation_reduce (param_0.370: f32[4096,3]) -> f32[] {
  %param_0.370 = f32[4096,3]{1,0} parameter(0)
  %multiply.233 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %param_0.370, f32[4096,3]{1,0} %param_0.370), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %bitcast.73 = f32[12288]{0} bitcast(f32[4096,3]{1,0} %multiply.233)
  %constant_472 = f32[] constant(0)
  ROOT %reduce.36 = f32[] reduce(f32[12288]{0} %bitcast.73, f32[] %constant_472), dimensions={0}, to_apply=%region_10.1472, metadata={op_name="jit(training_step)/jit(main)/reduce_sum[axes=(0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
}

%fused_computation.2 (param_0.12: u32[2], param_1.10: u32[2]) -> (u32[2], u32[2], u32[2]) {
  %param_0.12 = u32[2]{0} parameter(0)
  %param_1.10 = u32[2]{0} parameter(1)
  %concatenate.100 = u32[4]{0} concatenate(u32[2]{0} %param_0.12, u32[2]{0} %param_1.10), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %bitcast.75 = u32[2,2]{1,0} bitcast(u32[4]{0} %concatenate.100), metadata={op_name="jit(training_step)/jit(main)/reshape[new_sizes=(2, 2) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %slice.274 = u32[1,2]{1,0} slice(u32[2,2]{1,0} %bitcast.75), slice={[1:2], [0:2]}, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(1, 0) limit_indices=(2, 2) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %bitcast.74 = u32[2]{0} bitcast(u32[1,2]{1,0} %slice.274), metadata={op_name="jit(training_step)/jit(main)/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %slice.358.clone.1 = u32[1]{0} slice(u32[2]{0} %param_0.12), slice={[1:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(1,) limit_indices=(2,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %bitcast.133.clone.1 = u32[] bitcast(u32[1]{0} %slice.358.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %broadcast.954.clone.1 = u32[2]{0} broadcast(u32[] %bitcast.133.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %slice.360.clone.1 = u32[1]{0} slice(u32[2]{0} %param_0.12), slice={[0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(0,) limit_indices=(1,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %bitcast.134.clone.1 = u32[] bitcast(u32[1]{0} %slice.360.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %broadcast.955.clone.1 = u32[2]{0} broadcast(u32[] %bitcast.134.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  ROOT %tuple.87 = (u32[2]{0}, u32[2]{0}, u32[2]{0}) tuple(u32[2]{0} %bitcast.74, u32[2]{0} %broadcast.954.clone.1, u32[2]{0} %broadcast.955.clone.1)
}

%fused_computation.3 (param_0.13: f32[3,16,128,128], param_1.967: f32[], param_2.991: f32[], param_3.874: f32[], param_4.627: f32[3,16,128,128], param_5.567: f32[3,16,128,128], param_6.501: f32[3,16,128,128], param_7.433: f32[3,16,128,128]) -> (f32[3,16,128,128], f32[3,16,128,128], f32[3,16,128,128]) {
  %param_0.13 = f32[3,16,128,128]{3,2,1,0} parameter(0)
  %param_5.567 = f32[3,16,128,128]{3,2,1,0} parameter(5)
  %param_6.501 = f32[3,16,128,128]{3,2,1,0} parameter(6)
  %add.491.clone.1 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %param_5.567, f32[3,16,128,128]{3,2,1,0} %param_6.501), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_478_clone_1 = f32[] constant(0.0100000007)
  %broadcast.799.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_478_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.241.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %add.491.clone.1, f32[3,16,128,128]{3,2,1,0} %broadcast.799.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_7.433 = f32[3,16,128,128]{3,2,1,0} parameter(7)
  %constant_479_clone_1 = f32[] constant(0.9)
  %broadcast.798.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_479_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.240.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %param_7.433, f32[3,16,128,128]{3,2,1,0} %broadcast.798.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.336.clone.1 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %multiply.241.clone.1, f32[3,16,128,128]{3,2,1,0} %multiply.240.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.991 = f32[] parameter(2)
  %broadcast.794 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %param_2.991), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_477_clone_1 = f32[] constant(0.1)
  %broadcast.797.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_477_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.239.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %add.491.clone.1, f32[3,16,128,128]{3,2,1,0} %broadcast.797.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.238.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %multiply.239.clone.1, f32[3,16,128,128]{3,2,1,0} %multiply.239.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_475_clone_1 = f32[] constant(0.01)
  %broadcast.796.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_475_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.237.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %multiply.238.clone.1, f32[3,16,128,128]{3,2,1,0} %broadcast.796.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.627 = f32[3,16,128,128]{3,2,1,0} parameter(4)
  %constant_476_clone_1 = f32[] constant(0.99)
  %broadcast.795.clone.1 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_476_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.236.clone.1 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %param_4.627, f32[3,16,128,128]{3,2,1,0} %broadcast.795.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.335.clone.1 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %multiply.237.clone.1, f32[3,16,128,128]{3,2,1,0} %multiply.236.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.967 = f32[] parameter(1)
  %broadcast.793 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %param_1.967), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.24 = f32[3,16,128,128]{3,2,1,0} divide(f32[3,16,128,128]{3,2,1,0} %add.335.clone.1, f32[3,16,128,128]{3,2,1,0} %broadcast.793), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.0 = f32[3,16,128,128]{3,2,1,0} sqrt(f32[3,16,128,128]{3,2,1,0} %divide.24), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_474 = f32[] constant(1e-08)
  %broadcast.792 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_474), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.334 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %sqrt.0, f32[3,16,128,128]{3,2,1,0} %broadcast.792), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.235 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %broadcast.794, f32[3,16,128,128]{3,2,1,0} %add.334)
  %divide.23 = f32[3,16,128,128]{3,2,1,0} divide(f32[3,16,128,128]{3,2,1,0} %add.336.clone.1, f32[3,16,128,128]{3,2,1,0} %multiply.235), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.874 = f32[] parameter(3)
  %constant_473 = f32[] constant(-0.02)
  %multiply.397 = f32[] multiply(f32[] %param_3.874, f32[] %constant_473), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.791 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %multiply.397), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.234 = f32[3,16,128,128]{3,2,1,0} multiply(f32[3,16,128,128]{3,2,1,0} %divide.23, f32[3,16,128,128]{3,2,1,0} %broadcast.791), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.333 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %param_0.13, f32[3,16,128,128]{3,2,1,0} %multiply.234), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.46 = (f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}) tuple(f32[3,16,128,128]{3,2,1,0} %add.333, f32[3,16,128,128]{3,2,1,0} %add.335.clone.1, f32[3,16,128,128]{3,2,1,0} %add.336.clone.1)
}

%region_24.1622 (Arg_0.1623: f32[], Arg_1.1624: f32[]) -> f32[] {
  %Arg_0.1623 = f32[] parameter(0)
  %Arg_1.1624 = f32[] parameter(1)
  ROOT %add.1625 = f32[] add(f32[] %Arg_0.1623, f32[] %Arg_1.1624), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter (param_0.553: f32[3,4096,16,221], param_1.751: f32[3,4096,221], param_2.713: f32[3,4096,221]) -> f32[3,16,128,128] {
  %constant_480 = f32[] constant(0)
  %broadcast.962 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_480), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/broadcast_in_dim[shape=(3, 16, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.135 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1067 = s32[] constant(0)
  %broadcast.1483 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1067), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.713 = f32[3,4096,221]{2,1,0} parameter(2)
  %floor.159 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_2.713), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.98 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.159), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1066 = s32[] constant(127)
  %broadcast.1482 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1066), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.87 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1483, s32[3,4096,221]{2,1,0} %convert.98, s32[3,4096,221]{2,1,0} %broadcast.1482), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.213 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.87, s32[3,4096,221]{2,1,0} %broadcast.1483), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1065 = s32[] constant(128)
  %broadcast.1481 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1065), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.637 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.87, s32[3,4096,221]{2,1,0} %broadcast.1481), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.225 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.213, s32[3,4096,221]{2,1,0} %add.637, s32[3,4096,221]{2,1,0} %clamp.87), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.212 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.225), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.751 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.158 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.751), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.97 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.158), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.86 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1483, s32[3,4096,221]{2,1,0} %convert.97, s32[3,4096,221]{2,1,0} %broadcast.1482), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.212 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.86, s32[3,4096,221]{2,1,0} %broadcast.1483), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.636 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.86, s32[3,4096,221]{2,1,0} %broadcast.1481), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.224 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.212, s32[3,4096,221]{2,1,0} %add.636, s32[3,4096,221]{2,1,0} %clamp.86), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.211 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.224), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.174 = s32[3,4096,221,3]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.135, s32[3,4096,221,1]{2,1,0,3} %bitcast.212, s32[3,4096,221,1]{2,1,0,3} %bitcast.211), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.19 = s32[3,4096,221,3]{3,2,1,0} copy(s32[3,4096,221,3]{2,1,0,3} %concatenate.174), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1033 = f32[] constant(1)
  %broadcast.1459 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_1033), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.184 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_2.713, f32[3,4096,221]{2,1,0} %floor.159), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.183 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1459, f32[3,4096,221]{2,1,0} %subtract.184), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.182 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.751, f32[3,4096,221]{2,1,0} %floor.158), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.181 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1459, f32[3,4096,221]{2,1,0} %subtract.182), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.449 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.183, f32[3,4096,221]{2,1,0} %subtract.181), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1458 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.449), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.553 = f32[3,4096,16,221]{3,2,1,0} parameter(0)
  %multiply.242 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1458, f32[3,4096,16,221]{3,2,1,0} %param_0.553), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.42 = f32[3,16,128,128]{3,2,1,0} scatter(f32[3,16,128,128]{3,2,1,0} %broadcast.962, s32[3,4096,221,3]{3,2,1,0} %copy.19, f32[3,4096,16,221]{3,2,1,0} %multiply.242), update_window_dims={2}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_24.1622, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_23.1618 (Arg_0.1619: f32[], Arg_1.1620: f32[]) -> f32[] {
  %Arg_0.1619 = f32[] parameter(0)
  %Arg_1.1620 = f32[] parameter(1)
  ROOT %add.1621 = f32[] add(f32[] %Arg_0.1619, f32[] %Arg_1.1620), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.1 (param_0.550: f32[3,4096,16,221], param_1.743: f32[3,16,128,128], param_2.706: f32[3,16,128,128], param_3.569: f32[3,4096,221], param_4.322: f32[3,4096,221]) -> f32[3,16,128,128] {
  %param_1.743 = f32[3,16,128,128]{3,2,1,0} parameter(1)
  %param_2.706 = f32[3,16,128,128]{3,2,1,0} parameter(2)
  %add.337 = f32[3,16,128,128]{3,2,1,0} add(f32[3,16,128,128]{3,2,1,0} %param_1.743, f32[3,16,128,128]{3,2,1,0} %param_2.706), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.131 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1014 = s32[] constant(0)
  %broadcast.1451 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1014), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_4.322 = f32[3,4096,221]{2,1,0} parameter(4)
  %floor.143 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_4.322), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.90 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.143), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1010 = s32[] constant(127)
  %broadcast.1450 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1010), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.79 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1451, s32[3,4096,221]{2,1,0} %convert.90, s32[3,4096,221]{2,1,0} %broadcast.1450), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.205 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.79, s32[3,4096,221]{2,1,0} %broadcast.1451), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1008 = s32[] constant(128)
  %broadcast.1449 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1008), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.629 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.79, s32[3,4096,221]{2,1,0} %broadcast.1449), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.217 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.205, s32[3,4096,221]{2,1,0} %add.629, s32[3,4096,221]{2,1,0} %clamp.79), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.204 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.217), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_3.569 = f32[3,4096,221]{2,1,0} parameter(3)
  %floor.142 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_3.569), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.89 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.142), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1012 = s32[] constant(1)
  %broadcast.1447 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1012), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.628 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.89, s32[3,4096,221]{2,1,0} %broadcast.1447), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.78 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1451, s32[3,4096,221]{2,1,0} %add.628, s32[3,4096,221]{2,1,0} %broadcast.1450), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.204 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.78, s32[3,4096,221]{2,1,0} %broadcast.1451), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.627 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.78, s32[3,4096,221]{2,1,0} %broadcast.1449), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.216 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.204, s32[3,4096,221]{2,1,0} %add.627, s32[3,4096,221]{2,1,0} %clamp.78), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.203 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.216), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.170 = s32[3,4096,221,3]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.131, s32[3,4096,221,1]{3,2,1,0} %bitcast.204, s32[3,4096,221,1]{3,2,1,0} %bitcast.203), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_956 = f32[] constant(1)
  %broadcast.1411 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_956), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.168 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_4.322, f32[3,4096,221]{2,1,0} %floor.143), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.167 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1411, f32[3,4096,221]{2,1,0} %subtract.168), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.166 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_3.569, f32[3,4096,221]{2,1,0} %floor.142), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.444 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.167, f32[3,4096,221]{2,1,0} %subtract.166), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1410 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.444), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.550 = f32[3,4096,16,221]{3,2,1,0} parameter(0)
  %multiply.243 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1410, f32[3,4096,16,221]{3,2,1,0} %param_0.550), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.43 = f32[3,16,128,128]{3,2,1,0} scatter(f32[3,16,128,128]{3,2,1,0} %add.337, s32[3,4096,221,3]{3,2,1,0} %concatenate.170, f32[3,4096,16,221]{3,2,1,0} %multiply.243), update_window_dims={2}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_23.1618, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_22.1614 (Arg_0.1615: f32[], Arg_1.1616: f32[]) -> f32[] {
  %Arg_0.1615 = f32[] parameter(0)
  %Arg_1.1616 = f32[] parameter(1)
  ROOT %add.1617 = f32[] add(f32[] %Arg_0.1615, f32[] %Arg_1.1616), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.2 (param_0.535: f32[3,4096,16,221], param_1.727: f32[3,4096,221], param_2.688: f32[3,4096,221]) -> f32[3,16,128,128] {
  %constant_481 = f32[] constant(0)
  %broadcast.963 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_481), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/broadcast_in_dim[shape=(3, 16, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.127 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_946 = s32[] constant(0)
  %broadcast.1403 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_946), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.688 = f32[3,4096,221]{2,1,0} parameter(2)
  %floor.123 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_2.688), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.78 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.123), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_949 = s32[] constant(1)
  %broadcast.1402 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_949), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.608 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.78, s32[3,4096,221]{2,1,0} %broadcast.1402), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_945 = s32[] constant(127)
  %broadcast.1401 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_945), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.67 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1403, s32[3,4096,221]{2,1,0} %add.608, s32[3,4096,221]{2,1,0} %broadcast.1401), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.193 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.67, s32[3,4096,221]{2,1,0} %broadcast.1403), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_944 = s32[] constant(128)
  %broadcast.1400 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_944), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.607 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.67, s32[3,4096,221]{2,1,0} %broadcast.1400), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.205 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.193, s32[3,4096,221]{2,1,0} %add.607, s32[3,4096,221]{2,1,0} %clamp.67), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.196 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.205), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.727 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.122 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.727), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.77 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.122), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.66 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1403, s32[3,4096,221]{2,1,0} %convert.77, s32[3,4096,221]{2,1,0} %broadcast.1401), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.192 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.66, s32[3,4096,221]{2,1,0} %broadcast.1403), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.606 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.66, s32[3,4096,221]{2,1,0} %broadcast.1400), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.204 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.192, s32[3,4096,221]{2,1,0} %add.606, s32[3,4096,221]{2,1,0} %clamp.66), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.195 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.204), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.166 = s32[3,4096,221,3]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.127, s32[3,4096,221,1]{2,1,0,3} %bitcast.196, s32[3,4096,221,1]{2,1,0,3} %bitcast.195), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.20 = s32[3,4096,221,3]{3,2,1,0} copy(s32[3,4096,221,3]{2,1,0,3} %concatenate.166), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.156 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_2.688, f32[3,4096,221]{2,1,0} %floor.123), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_903 = f32[] constant(1)
  %broadcast.1356 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_903), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.155 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.727, f32[3,4096,221]{2,1,0} %floor.122), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.154 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1356, f32[3,4096,221]{2,1,0} %subtract.155), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.439 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.156, f32[3,4096,221]{2,1,0} %subtract.154), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1355 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.439), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.535 = f32[3,4096,16,221]{3,2,1,0} parameter(0)
  %multiply.244 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1355, f32[3,4096,16,221]{3,2,1,0} %param_0.535), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.44 = f32[3,16,128,128]{3,2,1,0} scatter(f32[3,16,128,128]{3,2,1,0} %broadcast.963, s32[3,4096,221,3]{3,2,1,0} %copy.20, f32[3,4096,16,221]{3,2,1,0} %multiply.244), update_window_dims={2}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_22.1614, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_21.1610 (Arg_0.1611: f32[], Arg_1.1612: f32[]) -> f32[] {
  %Arg_0.1611 = f32[] parameter(0)
  %Arg_1.1612 = f32[] parameter(1)
  ROOT %add.1613 = f32[] add(f32[] %Arg_0.1611, f32[] %Arg_1.1612), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.3 (param_0.532: f32[3,4096,16,221], param_1.709: f32[3,4096,221], param_2.668: f32[3,4096,221]) -> f32[3,16,128,128] {
  %constant_482 = f32[] constant(0)
  %broadcast.964 = f32[3,16,128,128]{3,2,1,0} broadcast(f32[] %constant_482), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/broadcast_in_dim[shape=(3, 16, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.123 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_893 = s32[] constant(0)
  %broadcast.1346 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_893), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.668 = f32[3,4096,221]{2,1,0} parameter(2)
  %floor.103 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_2.668), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.66 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.103), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_892 = s32[] constant(1)
  %broadcast.1345 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_892), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.586 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.66, s32[3,4096,221]{2,1,0} %broadcast.1345), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_891 = s32[] constant(127)
  %broadcast.1344 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_891), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.55 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1346, s32[3,4096,221]{2,1,0} %add.586, s32[3,4096,221]{2,1,0} %broadcast.1344), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.181 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.55, s32[3,4096,221]{2,1,0} %broadcast.1346), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_890 = s32[] constant(128)
  %broadcast.1343 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_890), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.585 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.55, s32[3,4096,221]{2,1,0} %broadcast.1343), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.193 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.181, s32[3,4096,221]{2,1,0} %add.585, s32[3,4096,221]{2,1,0} %clamp.55), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.184 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.193), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.709 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.102 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.709), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.65 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.102), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.584 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.65, s32[3,4096,221]{2,1,0} %broadcast.1345), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.54 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1346, s32[3,4096,221]{2,1,0} %add.584, s32[3,4096,221]{2,1,0} %broadcast.1344), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.180 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.54, s32[3,4096,221]{2,1,0} %broadcast.1346), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.582 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.54, s32[3,4096,221]{2,1,0} %broadcast.1343), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.192 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.180, s32[3,4096,221]{2,1,0} %add.582, s32[3,4096,221]{2,1,0} %clamp.54), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.183 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.192), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.162 = s32[3,4096,221,3]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.123, s32[3,4096,221,1]{3,2,1,0} %bitcast.184, s32[3,4096,221,1]{3,2,1,0} %bitcast.183), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.144 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_2.668, f32[3,4096,221]{2,1,0} %floor.103), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.143 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.709, f32[3,4096,221]{2,1,0} %floor.102), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.435 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.144, f32[3,4096,221]{2,1,0} %subtract.143), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1266 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.435), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.532 = f32[3,4096,16,221]{3,2,1,0} parameter(0)
  %multiply.245 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1266, f32[3,4096,16,221]{3,2,1,0} %param_0.532), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.45 = f32[3,16,128,128]{3,2,1,0} scatter(f32[3,16,128,128]{3,2,1,0} %broadcast.964, s32[3,4096,221,3]{3,2,1,0} %concatenate.162, f32[3,4096,16,221]{3,2,1,0} %multiply.245), update_window_dims={2}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_21.1610, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%fused_computation.6 (param_0.745: f32[3,16,128], param_1.965: f32[3,4096,221], param_2.989: f32[4096,221], param_3.872: f32[4096,221], param_4.625: pred[4096], param_5.564: f32[4096], param_6.494: f32[4096], param_7.424: f32[4096,221], param_8.326: f32[4096,221], param_9.233: f32[4096,221], param_10.137: f32[4096,110], param_11.102: f32[4096,110]) -> f32[3,4096,16,221] {
  %constant_803 = f32[] constant(1)
  %broadcast.1241 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_803), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.965 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.71 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.965), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.130 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.965, f32[3,4096,221]{2,1,0} %floor.71), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.129 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1241, f32[3,4096,221]{2,1,0} %subtract.130), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1240 = f32[3,4096,16,221]{2,1,0,3} broadcast(f32[3,4096,221]{2,1,0} %subtract.129), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.745 = f32[3,16,128]{2,1,0} parameter(0)
  %iota.115 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_799 = s32[] constant(0)
  %broadcast.1239 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_799), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.46 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.71), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_797 = s32[] constant(127)
  %broadcast.1238 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_797), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.35 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1239, s32[3,4096,221]{2,1,0} %convert.46, s32[3,4096,221]{2,1,0} %broadcast.1238), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.161 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.35, s32[3,4096,221]{2,1,0} %broadcast.1239), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_796 = s32[] constant(128)
  %broadcast.1237 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_796), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.536 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.35, s32[3,4096,221]{2,1,0} %broadcast.1237), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.173 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.161, s32[3,4096,221]{2,1,0} %add.536, s32[3,4096,221]{2,1,0} %clamp.35), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.168 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.173), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.154 = s32[3,4096,221,2]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.115, s32[3,4096,221,1]{2,1,0,3} %bitcast.168), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.83 = f32[3,4096,16,221]{2,1,0,3} gather(f32[3,16,128]{2,1,0} %param_0.745, s32[3,4096,221,2]{2,1,0,3} %concatenate.154), offset_dims={2}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,16,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 16, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.426 = f32[3,4096,16,221]{2,1,0,3} multiply(f32[3,4096,16,221]{2,1,0,3} %broadcast.1240, f32[3,4096,16,221]{2,1,0,3} %gather.83), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.33 = f32[3,4096,16,221]{3,2,1,0} copy(f32[3,4096,16,221]{2,1,0,3} %multiply.426), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1236 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %subtract.130), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.114 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_798 = s32[] constant(1)
  %broadcast.1234 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_798), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.535 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.46, s32[3,4096,221]{2,1,0} %broadcast.1234), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.34 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1239, s32[3,4096,221]{2,1,0} %add.535, s32[3,4096,221]{2,1,0} %broadcast.1238), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.160 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.34, s32[3,4096,221]{2,1,0} %broadcast.1239), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.534 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.34, s32[3,4096,221]{2,1,0} %broadcast.1237), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.172 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.160, s32[3,4096,221]{2,1,0} %add.534, s32[3,4096,221]{2,1,0} %clamp.34), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.167 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.172), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.153 = s32[3,4096,221,2]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.114, s32[3,4096,221,1]{3,2,1,0} %bitcast.167), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.82 = f32[3,4096,16,221]{3,2,1,0} gather(f32[3,16,128]{2,1,0} %param_0.745, s32[3,4096,221,2]{3,2,1,0} %concatenate.153), offset_dims={2}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,16,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 16, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.424 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1236, f32[3,4096,16,221]{3,2,1,0} %gather.82), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.533 = f32[3,4096,16,221]{3,2,1,0} add(f32[3,4096,16,221]{3,2,1,0} %copy.33, f32[3,4096,16,221]{3,2,1,0} %multiply.424), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_9.233 = f32[4096,221]{1,0} parameter(9)
  %slice.706 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_9.233), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_10.137 = f32[4096,110]{1,0} parameter(10)
  %param_11.102 = f32[4096,110]{1,0} parameter(11)
  %add.1015 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %param_10.137, f32[4096,110]{1,0} %param_11.102), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.705 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_9.233), slice={[0:4096:1], [0:219:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.29 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.705), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.1014 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %add.1015, f32[4096,110]{1,0} %reverse.29), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.263 = f32[4096,111]{1,0} concatenate(f32[4096,1]{1,0} %slice.706, f32[4096,110]{1,0} %add.1014), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_1803 = f32[] constant(0)
  %pad.176 = f32[4096,221]{1,0} pad(f32[4096,111]{1,0} %concatenate.263, f32[] %constant_1803), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.175 = f32[4096,221]{1,0} pad(f32[4096,110]{1,0} %add.1015, f32[] %constant_1803), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.1013 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %pad.176, f32[4096,221]{1,0} %pad.175), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.28 = f32[4096,221]{1,0} reverse(f32[4096,221]{1,0} %add.1013), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/rev[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_7.424 = f32[4096,221]{1,0} parameter(7)
  %param_8.326 = f32[4096,221]{1,0} parameter(8)
  %slice.704 = f32[4096,220]{1,0} slice(f32[4096,221]{1,0} %param_8.326), slice={[0:4096], [0:220]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(4096, 220) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=286}
  %pad.174 = f32[4096,221]{1,0} pad(f32[4096,220]{1,0} %slice.704, f32[] %constant_803), padding=0_0x1_0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=283}
  %multiply.617 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %param_7.424, f32[4096,221]{1,0} %pad.174), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=280}
  %negate.55 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.617), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %param_2.989 = f32[4096,221]{1,0} parameter(2)
  %negate.54 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_2.989), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_4.625 = pred[4096]{0} parameter(4)
  %param_6.494 = f32[4096]{0} parameter(6)
  %param_5.564 = f32[4096]{0} parameter(5)
  %constant_1804 = f32[] constant(0.001)
  %broadcast.2270 = f32[4096]{0} broadcast(f32[] %constant_1804), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.1012 = f32[4096]{0} add(f32[4096]{0} %param_5.564, f32[4096]{0} %broadcast.2270), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.97 = f32[4096]{0} maximum(f32[4096]{0} %param_6.494, f32[4096]{0} %add.1012), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.419 = f32[4096]{0} select(pred[4096]{0} %param_4.625, f32[4096]{0} %maximum.97, f32[4096]{0} %broadcast.2270), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %broadcast.2269 = f32[4096]{0} broadcast(f32[] %constant_1803), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.417 = f32[4096]{0} select(pred[4096]{0} %param_4.625, f32[4096]{0} %param_5.564, f32[4096]{0} %broadcast.2269), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.390 = f32[4096]{0} subtract(f32[4096]{0} %select.419, f32[4096]{0} %select.417), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1801 = f32[] constant(0.00452488707)
  %broadcast.2268 = f32[4096]{0} broadcast(f32[] %constant_1801), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.616 = f32[4096]{0} multiply(f32[4096]{0} %subtract.390, f32[4096]{0} %broadcast.2268), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.2267 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.616), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.615 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.54, f32[4096,221]{1,0} %broadcast.2267), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %exponential.32 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %multiply.615), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.614 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.55, f32[4096,221]{1,0} %exponential.32), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %add.1011 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %reverse.28, f32[4096,221]{1,0} %multiply.614), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.612 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %add.1011, f32[4096,221]{1,0} %broadcast.2267), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %negate.53 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.612), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_3.872 = f32[4096,221]{1,0} parameter(3)
  %constant_1800 = f32[] constant(inf)
  %broadcast.2262 = f32[4096,221]{1,0} broadcast(f32[] %constant_1800), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.351 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %param_3.872, f32[4096,221]{1,0} %broadcast.2262), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2261 = f32[4096,221]{1,0} broadcast(f32[] %constant_1803), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.413 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.351, f32[4096,221]{1,0} %broadcast.2261, f32[4096,221]{1,0} %param_3.872), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.350 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %param_2.989, f32[4096,221]{1,0} %broadcast.2262), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.412 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.350, f32[4096,221]{1,0} %broadcast.2261, f32[4096,221]{1,0} %param_2.989), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %subtract.388 = f32[4096,221]{1,0} subtract(f32[4096,221]{1,0} %select.413, f32[4096,221]{1,0} %select.412), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %exponential.31 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %subtract.388), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %multiply.611 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.53, f32[4096,221]{1,0} %exponential.31), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(softplus)))/jit(jit_transpose(jvp(softplus)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2260 = f32[3,16,4096,221]{1,2,0,3} broadcast(f32[4096,221]{1,0} %multiply.611), dimensions={2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 16, 4096, 221) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %transpose.51 = f32[3,4096,16,221]{2,1,0,3} transpose(f32[3,16,4096,221]{1,2,0,3} %broadcast.2260), dimensions={0,2,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/transpose[permutation=(0, 2, 1, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=90}
  %copy.21 = f32[3,4096,16,221]{3,2,1,0} copy(f32[3,4096,16,221]{2,1,0,3} %transpose.51), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/transpose[permutation=(0, 2, 1, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=90}
  ROOT %multiply.246 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %add.533, f32[3,4096,16,221]{3,2,1,0} %copy.21), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
}

%fused_computation.7 (param_0.24: f32[3,16,128], param_1.968: f32[], param_2.993: f32[], param_3.876: f32[], param_4.629: f32[3,16,128], param_5.570: f32[3,16,128], param_6.508: f32[3,16,128], param_7.442: f32[3,16,128]) -> (f32[3,16,128], f32[3,16,128], f32[3,16,128]) {
  %param_0.24 = f32[3,16,128]{2,1,0} parameter(0)
  %param_5.570 = f32[3,16,128]{2,1,0} parameter(5)
  %param_6.508 = f32[3,16,128]{2,1,0} parameter(6)
  %add.493.clone.1 = f32[3,16,128]{2,1,0} add(f32[3,16,128]{2,1,0} %param_5.570, f32[3,16,128]{2,1,0} %param_6.508), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_488_clone_1 = f32[] constant(0.0100000007)
  %broadcast.811.clone.1 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_488_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.254.clone.1 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %add.493.clone.1, f32[3,16,128]{2,1,0} %broadcast.811.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_7.442 = f32[3,16,128]{2,1,0} parameter(7)
  %constant_489_clone_1 = f32[] constant(0.9)
  %broadcast.810.clone.1 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_489_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.253.clone.1 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %param_7.442, f32[3,16,128]{2,1,0} %broadcast.810.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.341.clone.1 = f32[3,16,128]{2,1,0} add(f32[3,16,128]{2,1,0} %multiply.254.clone.1, f32[3,16,128]{2,1,0} %multiply.253.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.993 = f32[] parameter(2)
  %broadcast.804 = f32[3,16,128]{2,1,0} broadcast(f32[] %param_2.993), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_487_clone_1 = f32[] constant(0.1)
  %broadcast.808.clone.1 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_487_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.252.clone.1 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %add.493.clone.1, f32[3,16,128]{2,1,0} %broadcast.808.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.251.clone.1 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %multiply.252.clone.1, f32[3,16,128]{2,1,0} %multiply.252.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_485_clone_1 = f32[] constant(0.01)
  %broadcast.806.clone.1 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_485_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.250.clone.1 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %multiply.251.clone.1, f32[3,16,128]{2,1,0} %broadcast.806.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.629 = f32[3,16,128]{2,1,0} parameter(4)
  %constant_486_clone_1 = f32[] constant(0.99)
  %broadcast.805.clone.1 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_486_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.249.clone.1 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %param_4.629, f32[3,16,128]{2,1,0} %broadcast.805.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.340.clone.1 = f32[3,16,128]{2,1,0} add(f32[3,16,128]{2,1,0} %multiply.250.clone.1, f32[3,16,128]{2,1,0} %multiply.249.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.968 = f32[] parameter(1)
  %broadcast.803 = f32[3,16,128]{2,1,0} broadcast(f32[] %param_1.968), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.26 = f32[3,16,128]{2,1,0} divide(f32[3,16,128]{2,1,0} %add.340.clone.1, f32[3,16,128]{2,1,0} %broadcast.803), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.1 = f32[3,16,128]{2,1,0} sqrt(f32[3,16,128]{2,1,0} %divide.26), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_484 = f32[] constant(1e-08)
  %broadcast.801 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_484), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.339 = f32[3,16,128]{2,1,0} add(f32[3,16,128]{2,1,0} %sqrt.1, f32[3,16,128]{2,1,0} %broadcast.801), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.248 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %broadcast.804, f32[3,16,128]{2,1,0} %add.339)
  %divide.25 = f32[3,16,128]{2,1,0} divide(f32[3,16,128]{2,1,0} %add.341.clone.1, f32[3,16,128]{2,1,0} %multiply.248), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.876 = f32[] parameter(3)
  %constant_483 = f32[] constant(-0.02)
  %multiply.398 = f32[] multiply(f32[] %param_3.876, f32[] %constant_483), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.800 = f32[3,16,128]{2,1,0} broadcast(f32[] %multiply.398), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.247 = f32[3,16,128]{2,1,0} multiply(f32[3,16,128]{2,1,0} %divide.25, f32[3,16,128]{2,1,0} %broadcast.800), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.338 = f32[3,16,128]{2,1,0} add(f32[3,16,128]{2,1,0} %param_0.24, f32[3,16,128]{2,1,0} %multiply.247), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.48 = (f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}) tuple(f32[3,16,128]{2,1,0} %add.338, f32[3,16,128]{2,1,0} %add.340.clone.1, f32[3,16,128]{2,1,0} %add.341.clone.1)
}

%region_26.1665 (Arg_0.1666: f32[], Arg_1.1667: f32[]) -> f32[] {
  %Arg_0.1666 = f32[] parameter(0)
  %Arg_1.1667 = f32[] parameter(1)
  ROOT %add.1668 = f32[] add(f32[] %Arg_0.1666, f32[] %Arg_1.1667), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.4 (param_0.515: f32[3,4096,16,221], param_1.677: f32[3,4096,221]) -> f32[3,16,128] {
  %constant_490 = f32[] constant(0)
  %broadcast.965 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_490), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/broadcast_in_dim[shape=(3, 16, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.111 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_786 = s32[] constant(0)
  %broadcast.1219 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_786), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.677 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.63 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.677), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.42 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.63), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_785 = s32[] constant(127)
  %broadcast.1218 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_785), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.31 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1219, s32[3,4096,221]{2,1,0} %convert.42, s32[3,4096,221]{2,1,0} %broadcast.1218), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.157 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.31, s32[3,4096,221]{2,1,0} %broadcast.1219), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_784 = s32[] constant(128)
  %broadcast.1217 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_784), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.527 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.31, s32[3,4096,221]{2,1,0} %broadcast.1217), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.169 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.157, s32[3,4096,221]{2,1,0} %add.527, s32[3,4096,221]{2,1,0} %clamp.31), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.164 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.169), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.150 = s32[3,4096,221,2]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.111, s32[3,4096,221,1]{2,1,0,3} %bitcast.164), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_768 = f32[] constant(1)
  %broadcast.1205 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_768), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.124 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.677, f32[3,4096,221]{2,1,0} %floor.63), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.123 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1205, f32[3,4096,221]{2,1,0} %subtract.124), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1204 = f32[3,4096,16,221]{2,1,0,3} broadcast(f32[3,4096,221]{2,1,0} %subtract.123), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.515 = f32[3,4096,16,221]{2,1,0,3} parameter(0)
  %multiply.255 = f32[3,4096,16,221]{2,1,0,3} multiply(f32[3,4096,16,221]{2,1,0,3} %broadcast.1204, f32[3,4096,16,221]{2,1,0,3} %param_0.515), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.46 = f32[3,16,128]{2,1,0} scatter(f32[3,16,128]{2,1,0} %broadcast.965, s32[3,4096,221,2]{2,1,0,3} %concatenate.150, f32[3,4096,16,221]{2,1,0,3} %multiply.255), update_window_dims={2}, inserted_window_dims={0,2}, scatter_dims_to_operand_dims={0,2}, index_vector_dim=3, to_apply=%region_26.1665, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_25.1661 (Arg_0.1662: f32[], Arg_1.1663: f32[]) -> f32[] {
  %Arg_0.1662 = f32[] parameter(0)
  %Arg_1.1663 = f32[] parameter(1)
  ROOT %add.1664 = f32[] add(f32[] %Arg_0.1662, f32[] %Arg_1.1663), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.5 (param_0.510: f32[3,4096,16,221], param_1.673: f32[3,4096,221]) -> f32[3,16,128] {
  %constant_491 = f32[] constant(0)
  %broadcast.966 = f32[3,16,128]{2,1,0} broadcast(f32[] %constant_491), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/broadcast_in_dim[shape=(3, 16, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.107 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_760 = s32[] constant(0)
  %broadcast.1197 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_760), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.673 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.55 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.673), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.38 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.55), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_759 = s32[] constant(1)
  %broadcast.1196 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_759), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.522 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.38, s32[3,4096,221]{2,1,0} %broadcast.1196), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_758 = s32[] constant(127)
  %broadcast.1195 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_758), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.27 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1197, s32[3,4096,221]{2,1,0} %add.522, s32[3,4096,221]{2,1,0} %broadcast.1195), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.153 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.27, s32[3,4096,221]{2,1,0} %broadcast.1197), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_757 = s32[] constant(128)
  %broadcast.1193 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_757), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.521 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.27, s32[3,4096,221]{2,1,0} %broadcast.1193), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.165 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.153, s32[3,4096,221]{2,1,0} %add.521, s32[3,4096,221]{2,1,0} %clamp.27), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.160 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.165), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.146 = s32[3,4096,221,2]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.107, s32[3,4096,221,1]{3,2,1,0} %bitcast.160), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.77 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.673, f32[3,4096,221]{2,1,0} %floor.55), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.967 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %subtract.77), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.510 = f32[3,4096,16,221]{2,1,0,3} parameter(0)
  %copy.22 = f32[3,4096,16,221]{3,2,1,0} copy(f32[3,4096,16,221]{2,1,0,3} %param_0.510), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %multiply.256 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.967, f32[3,4096,16,221]{3,2,1,0} %copy.22), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.47 = f32[3,16,128]{2,1,0} scatter(f32[3,16,128]{2,1,0} %broadcast.966, s32[3,4096,221,2]{3,2,1,0} %concatenate.146, f32[3,4096,16,221]{3,2,1,0} %multiply.256), update_window_dims={2}, inserted_window_dims={0,2}, scatter_dims_to_operand_dims={0,2}, index_vector_dim=3, to_apply=%region_25.1661, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%fused_computation.10 (param_0.743: f32[3,16,128,128], param_1.963: f32[3,4096,221], param_2.984: f32[3,4096,221], param_3.865: f32[4096,221], param_4.616: f32[4096,221], param_5.553: pred[4096], param_6.480: f32[4096], param_7.414: f32[4096], param_8.318: f32[4096,221], param_9.219: f32[4096,221], param_10.123: f32[4096,221], param_11.85: f32[4096,110], param_12.47: f32[4096,110]) -> f32[3,4096,16,221] {
  %param_10.123 = f32[4096,221]{1,0} parameter(10)
  %slice.700 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_10.123), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_11.85 = f32[4096,110]{1,0} parameter(11)
  %param_12.47 = f32[4096,110]{1,0} parameter(12)
  %add.1003 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %param_11.85, f32[4096,110]{1,0} %param_12.47), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.699 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_10.123), slice={[0:4096:1], [0:219:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.25 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.699), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.1002 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %add.1003, f32[4096,110]{1,0} %reverse.25), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.261 = f32[4096,111]{1,0} concatenate(f32[4096,1]{1,0} %slice.700, f32[4096,110]{1,0} %add.1002), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_1783 = f32[] constant(0)
  %pad.170 = f32[4096,221]{1,0} pad(f32[4096,111]{1,0} %concatenate.261, f32[] %constant_1783), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.169 = f32[4096,221]{1,0} pad(f32[4096,110]{1,0} %add.1003, f32[] %constant_1783), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.1001 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %pad.170, f32[4096,221]{1,0} %pad.169), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.24 = f32[4096,221]{1,0} reverse(f32[4096,221]{1,0} %add.1001), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/rev[dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_8.318 = f32[4096,221]{1,0} parameter(8)
  %param_9.219 = f32[4096,221]{1,0} parameter(9)
  %slice.698 = f32[4096,220]{1,0} slice(f32[4096,221]{1,0} %param_9.219), slice={[0:4096], [0:220]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(4096, 220) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=286}
  %constant_1118 = f32[] constant(1)
  %pad.168 = f32[4096,221]{1,0} pad(f32[4096,220]{1,0} %slice.698, f32[] %constant_1118), padding=0_0x1_0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=283}
  %multiply.603 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %param_8.318, f32[4096,221]{1,0} %pad.168), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=280}
  %negate.49 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.603), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %param_3.865 = f32[4096,221]{1,0} parameter(3)
  %negate.48 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_3.865), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_5.553 = pred[4096]{0} parameter(5)
  %param_7.414 = f32[4096]{0} parameter(7)
  %param_6.480 = f32[4096]{0} parameter(6)
  %constant_1784 = f32[] constant(0.001)
  %broadcast.2243 = f32[4096]{0} broadcast(f32[] %constant_1784), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.1000 = f32[4096]{0} add(f32[4096]{0} %param_6.480, f32[4096]{0} %broadcast.2243), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.93 = f32[4096]{0} maximum(f32[4096]{0} %param_7.414, f32[4096]{0} %add.1000), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.405 = f32[4096]{0} select(pred[4096]{0} %param_5.553, f32[4096]{0} %maximum.93, f32[4096]{0} %broadcast.2243), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %broadcast.2241 = f32[4096]{0} broadcast(f32[] %constant_1783), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.404 = f32[4096]{0} select(pred[4096]{0} %param_5.553, f32[4096]{0} %param_6.480, f32[4096]{0} %broadcast.2241), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.384 = f32[4096]{0} subtract(f32[4096]{0} %select.405, f32[4096]{0} %select.404), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1782 = f32[] constant(0.00452488707)
  %broadcast.2239 = f32[4096]{0} broadcast(f32[] %constant_1782), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.602 = f32[4096]{0} multiply(f32[4096]{0} %subtract.384, f32[4096]{0} %broadcast.2239), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.2237 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.602), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.601 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.48, f32[4096,221]{1,0} %broadcast.2237), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %exponential.28 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %multiply.601), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.600 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.49, f32[4096,221]{1,0} %exponential.28), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %add.999 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %reverse.24, f32[4096,221]{1,0} %multiply.600), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %multiply.598 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %add.999, f32[4096,221]{1,0} %broadcast.2237), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %negate.47 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %multiply.598), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_4.616 = f32[4096,221]{1,0} parameter(4)
  %constant_1781 = f32[] constant(inf)
  %broadcast.2229 = f32[4096,221]{1,0} broadcast(f32[] %constant_1781), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.347 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %param_4.616, f32[4096,221]{1,0} %broadcast.2229), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2228 = f32[4096,221]{1,0} broadcast(f32[] %constant_1783), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.401 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.347, f32[4096,221]{1,0} %broadcast.2228, f32[4096,221]{1,0} %param_4.616), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %compare.346 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %param_3.865, f32[4096,221]{1,0} %broadcast.2229), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %select.399 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.346, f32[4096,221]{1,0} %broadcast.2228, f32[4096,221]{1,0} %param_3.865), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %subtract.382 = f32[4096,221]{1,0} subtract(f32[4096,221]{1,0} %select.401, f32[4096,221]{1,0} %select.399), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %exponential.27 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %subtract.382), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %multiply.597 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.47, f32[4096,221]{1,0} %exponential.27), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(softplus)))/jit(jit_transpose(jvp(softplus)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %broadcast.2227 = f32[3,16,4096,221]{1,2,0,3} broadcast(f32[4096,221]{1,0} %multiply.597), dimensions={2,3}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 16, 4096, 221) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %transpose.49 = f32[3,4096,16,221]{2,1,0,3} transpose(f32[3,16,4096,221]{1,2,0,3} %broadcast.2227), dimensions={0,2,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/transpose[permutation=(0, 2, 1, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=90}
  %broadcast.1553 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_1118), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.984 = f32[3,4096,221]{2,1,0} parameter(2)
  %floor.191 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_2.984), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.208 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_2.984, f32[3,4096,221]{2,1,0} %floor.191), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.207 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1553, f32[3,4096,221]{2,1,0} %subtract.208), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.963 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.190 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.963), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.206 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.963, f32[3,4096,221]{2,1,0} %floor.190), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.205 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1553, f32[3,4096,221]{2,1,0} %subtract.206), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.465 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.207, f32[3,4096,221]{2,1,0} %subtract.205), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1552 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.465), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.743 = f32[3,16,128,128]{3,2,1,0} parameter(0)
  %iota.143 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1106 = s32[] constant(0)
  %broadcast.1551 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1106), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.114 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.191), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1104 = s32[] constant(127)
  %broadcast.1550 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1104), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.103 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1551, s32[3,4096,221]{2,1,0} %convert.114, s32[3,4096,221]{2,1,0} %broadcast.1550), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.229 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.103, s32[3,4096,221]{2,1,0} %broadcast.1551), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1103 = s32[] constant(128)
  %broadcast.1549 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1103), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.667 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.103, s32[3,4096,221]{2,1,0} %broadcast.1549), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.241 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.229, s32[3,4096,221]{2,1,0} %add.667, s32[3,4096,221]{2,1,0} %clamp.103), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.228 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.241), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.113 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.190), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.102 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1551, s32[3,4096,221]{2,1,0} %convert.113, s32[3,4096,221]{2,1,0} %broadcast.1550), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.228 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.102, s32[3,4096,221]{2,1,0} %broadcast.1551), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.666 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.102, s32[3,4096,221]{2,1,0} %broadcast.1549), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.240 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.228, s32[3,4096,221]{2,1,0} %add.666, s32[3,4096,221]{2,1,0} %clamp.102), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.227 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.240), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.182 = s32[3,4096,221,3]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.143, s32[3,4096,221,1]{2,1,0,3} %bitcast.228, s32[3,4096,221,1]{2,1,0,3} %bitcast.227), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.95 = f32[3,4096,16,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_0.743, s32[3,4096,221,3]{2,1,0,3} %concatenate.182), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.464 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1552, f32[3,4096,16,221]{3,2,1,0} %gather.95), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.463 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.207, f32[3,4096,221]{2,1,0} %subtract.206), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1544 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.463), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.142 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.226 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.241), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1105 = s32[] constant(1)
  %broadcast.1539 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_1105), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.664 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.113, s32[3,4096,221]{2,1,0} %broadcast.1539), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.100 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1551, s32[3,4096,221]{2,1,0} %add.664, s32[3,4096,221]{2,1,0} %broadcast.1550), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.226 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.100, s32[3,4096,221]{2,1,0} %broadcast.1551), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.663 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.100, s32[3,4096,221]{2,1,0} %broadcast.1549), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.238 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.226, s32[3,4096,221]{2,1,0} %add.663, s32[3,4096,221]{2,1,0} %clamp.100), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.225 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.238), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.181 = s32[3,4096,221,3]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.142, s32[3,4096,221,1]{3,2,1,0} %bitcast.226, s32[3,4096,221,1]{3,2,1,0} %bitcast.225), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.94 = f32[3,4096,16,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_0.743, s32[3,4096,221,3]{3,2,1,0} %concatenate.181), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.462 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1544, f32[3,4096,16,221]{3,2,1,0} %gather.94), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.662 = f32[3,4096,16,221]{3,2,1,0} add(f32[3,4096,16,221]{3,2,1,0} %multiply.464, f32[3,4096,16,221]{3,2,1,0} %multiply.462), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.461 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.208, f32[3,4096,221]{2,1,0} %subtract.205), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1535 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.461), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.141 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.661 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.114, s32[3,4096,221]{2,1,0} %broadcast.1539), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.99 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1551, s32[3,4096,221]{2,1,0} %add.661, s32[3,4096,221]{2,1,0} %broadcast.1550), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.225 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.99, s32[3,4096,221]{2,1,0} %broadcast.1551), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.660 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.99, s32[3,4096,221]{2,1,0} %broadcast.1549), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.237 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.225, s32[3,4096,221]{2,1,0} %add.660, s32[3,4096,221]{2,1,0} %clamp.99), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.224 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.237), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.180 = s32[3,4096,221,3]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.141, s32[3,4096,221,1]{2,1,0,3} %bitcast.224, s32[3,4096,221,1]{2,1,0,3} %bitcast.227), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.93 = f32[3,4096,16,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_0.743, s32[3,4096,221,3]{2,1,0,3} %concatenate.180), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.460 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1535, f32[3,4096,16,221]{3,2,1,0} %gather.93), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.658 = f32[3,4096,16,221]{3,2,1,0} add(f32[3,4096,16,221]{3,2,1,0} %add.662, f32[3,4096,16,221]{3,2,1,0} %multiply.460), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.459 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.208, f32[3,4096,221]{2,1,0} %subtract.206), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1527 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.459), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.140 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.222 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.237), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.179 = s32[3,4096,221,3]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.140, s32[3,4096,221,1]{3,2,1,0} %bitcast.222, s32[3,4096,221,1]{3,2,1,0} %bitcast.225), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.92 = f32[3,4096,16,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_0.743, s32[3,4096,221,3]{3,2,1,0} %concatenate.179), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.458 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1527, f32[3,4096,16,221]{3,2,1,0} %gather.92), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.653 = f32[3,4096,16,221]{3,2,1,0} add(f32[3,4096,16,221]{3,2,1,0} %add.658, f32[3,4096,16,221]{3,2,1,0} %multiply.458), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.23 = f32[3,4096,16,221]{2,1,0,3} copy(f32[3,4096,16,221]{3,2,1,0} %add.653), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %multiply.257 = f32[3,4096,16,221]{2,1,0,3} multiply(f32[3,4096,16,221]{2,1,0,3} %transpose.49, f32[3,4096,16,221]{2,1,0,3} %copy.23), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
}

%fused_computation.16 (param_0.728: f32[4096,3], param_1.937: f32[4096,13]) -> f32[4096,6] {
  %param_1.937 = f32[4096,13]{1,0} parameter(1)
  %slice.674 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_1.937), slice={[0:4096:1], [0:12:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 12) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.673 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_1.937), slice={[0:4096:1], [1:13:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 13) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.965 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %slice.674, f32[4096,6]{1,0} %slice.673), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.288 = f32[4096,1]{1,0} slice(f32[4096,6]{1,0} %add.965), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_0.728 = f32[4096,3]{1,0} parameter(0)
  %slice.287 = f32[4096,2]{1,0} slice(f32[4096,3]{1,0} %param_0.728), slice={[0:4096], [0:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 2) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.286 = f32[4096,2]{1,0} slice(f32[4096,6]{1,0} %add.965), slice={[0:4096:1], [2:6:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 6) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.354 = f32[4096,2]{1,0} add(f32[4096,2]{1,0} %slice.287, f32[4096,2]{1,0} %slice.286), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.106 = f32[4096,3]{1,0} concatenate(f32[4096,1]{1,0} %slice.288, f32[4096,2]{1,0} %add.354), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_498 = f32[] constant(0)
  ROOT %pad.104 = f32[4096,6]{1,0} pad(f32[4096,3]{1,0} %concatenate.106, f32[] %constant_498), padding=0_0_0x0_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.17 (param_0.731: f32[4096,13]) -> f32[4096,3] {
  %param_0.731 = f32[4096,13]{1,0} parameter(0)
  %slice.678 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_0.731), slice={[0:4096:1], [0:12:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 12) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.677 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_0.731), slice={[0:4096:1], [1:13:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 13) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.967 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %slice.678, f32[4096,6]{1,0} %slice.677), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.294 = f32[4096,3]{1,0} slice(f32[4096,6]{1,0} %add.967), slice={[0:4096:1], [0:5:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 5) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.293 = f32[4096,3]{1,0} slice(f32[4096,6]{1,0} %add.967), slice={[0:4096:1], [1:6:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 6) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.358 = f32[4096,3]{1,0} add(f32[4096,3]{1,0} %slice.294, f32[4096,3]{1,0} %slice.293), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.292 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.358), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.291 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.358), slice={[0:4096:1], [0:2:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 2) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.290 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.358), slice={[0:4096:1], [1:3:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 3) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.357 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %slice.291, f32[4096,1]{1,0} %slice.290), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.289 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.358), slice={[0:4096:1], [2:3:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 3) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.356 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %add.357, f32[4096,1]{1,0} %slice.289), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.107 = f32[4096,2]{1,0} concatenate(f32[4096,1]{1,0} %slice.292, f32[4096,1]{1,0} %add.356), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_499 = f32[] constant(0)
  %pad.106 = f32[4096,3]{1,0} pad(f32[4096,2]{1,0} %concatenate.107, f32[] %constant_499), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.105 = f32[4096,3]{1,0} pad(f32[4096,1]{1,0} %add.357, f32[] %constant_499), padding=0_0x1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %add.355 = f32[4096,3]{1,0} add(f32[4096,3]{1,0} %pad.106, f32[4096,3]{1,0} %pad.105), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.19 (param_0.726: f32[4096,55]) -> f32[4096,13] {
  %param_0.726 = f32[4096,55]{1,0} parameter(0)
  %slice.670 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_0.726), slice={[0:4096:1], [0:54:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 54) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.669 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_0.726), slice={[0:4096:1], [1:55:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 55) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.963 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %slice.670, f32[4096,27]{1,0} %slice.669), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.298 = f32[4096,13]{1,0} slice(f32[4096,27]{1,0} %add.963), slice={[0:4096:1], [0:26:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 26) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.297 = f32[4096,13]{1,0} slice(f32[4096,27]{1,0} %add.963), slice={[0:4096:1], [1:27:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 27) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %add.360 = f32[4096,13]{1,0} add(f32[4096,13]{1,0} %slice.298, f32[4096,13]{1,0} %slice.297), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.21 (param_0.722: f32[4096,221]) -> f32[4096,55] {
  %param_0.722 = f32[4096,221]{1,0} parameter(0)
  %slice.662 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_0.722), slice={[0:4096:1], [2:221:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 220) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.21 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.662), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 220) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.661 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_0.722), slice={[0:4096:1], [1:220:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.20 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.661), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.958 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %reverse.21, f32[4096,110]{1,0} %reverse.20), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.302 = f32[4096,55]{1,0} slice(f32[4096,110]{1,0} %add.958), slice={[0:4096:1], [0:109:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 109) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.301 = f32[4096,55]{1,0} slice(f32[4096,110]{1,0} %add.958), slice={[0:4096:1], [1:110:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 110) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %add.362 = f32[4096,55]{1,0} add(f32[4096,55]{1,0} %slice.302, f32[4096,55]{1,0} %slice.301), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%region_11.1484 (Arg_0.1485: f32[], Arg_1.1486: f32[]) -> f32[] {
  %Arg_0.1485 = f32[] parameter(0)
  %Arg_1.1486 = f32[] parameter(1)
  ROOT %add.1487 = f32[] add(f32[] %Arg_0.1485, f32[] %Arg_1.1486), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=196}
}

%fused_computation.23 (param_0.62: f32[4096,221], param_1.779: f32[4096,221], param_2.747: f32[4096], param_3.832: pred[4096], param_4.605: f32[4096], param_5.544: f32[4096], param_6.469: f32[4096,221], param_7.400: f32[4096], param_8.308: f32[4096,3]) -> f32[4096,221] {
  %param_8.308 = f32[4096,3]{1,0} parameter(8)
  %constant_1751 = f32[] constant(0.000162760422)
  %broadcast.2215 = f32[4096,3]{1,0} broadcast(f32[] %constant_1751), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %multiply.589 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %param_8.308, f32[4096,3]{1,0} %broadcast.2215), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %constant_500 = f32[] constant(0)
  %reduce.37 = f32[4096]{0} reduce(f32[4096,3]{1,0} %multiply.589, f32[] %constant_500), dimensions={1}, to_apply=%region_11.1484, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=196}
  %bitcast.77 = f32[4096,1]{1,0} bitcast(f32[4096]{0} %reduce.37)
  %pad.109 = f32[4096,221]{1,0} pad(f32[4096,1]{1,0} %bitcast.77, f32[] %constant_500), padding=0_0x220_0, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0, 1), inserted_window_dims=(), scatter_dims_to_operand_dims=(1,)) indices_are_sorted=True unique_indices=True mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=196}
  %param_2.747 = f32[4096]{0} parameter(2)
  %param_7.400 = f32[4096]{0} parameter(7)
  %constant_1737 = f32[] constant(1e-08)
  %broadcast.2199 = f32[4096]{0} broadcast(f32[] %constant_1737), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.953 = f32[4096]{0} add(f32[4096]{0} %param_7.400, f32[4096]{0} %broadcast.2199), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  %divide.27 = f32[4096]{0} divide(f32[4096]{0} %param_2.747, f32[4096]{0} %add.953), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %negate.18 = f32[4096]{0} negate(f32[4096]{0} %divide.27), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %bitcast.76 = f32[4096,1]{1,0} bitcast(f32[4096]{0} %negate.18)
  %pad.108 = f32[4096,221]{1,0} pad(f32[4096,1]{1,0} %bitcast.76, f32[] %constant_500), padding=0_0x220_0, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0,), inserted_window_dims=(1,), scatter_dims_to_operand_dims=(1,)) indices_are_sorted=True unique_indices=True mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=482}
  %add.365 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %pad.109, f32[4096,221]{1,0} %pad.108), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=482}
  %constant_1242 = f32[] constant(1)
  %broadcast.1689 = f32[4096,221]{1,0} broadcast(f32[] %constant_1242), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %param_6.469 = f32[4096,221]{1,0} parameter(6)
  %negate.35 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_6.469), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_3.832 = pred[4096]{0} parameter(3)
  %param_5.544 = f32[4096]{0} parameter(5)
  %param_4.605 = f32[4096]{0} parameter(4)
  %constant_1241 = f32[] constant(0.001)
  %broadcast.1688 = f32[4096]{0} broadcast(f32[] %constant_1241), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.711 = f32[4096]{0} add(f32[4096]{0} %param_4.605, f32[4096]{0} %broadcast.1688), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.85 = f32[4096]{0} maximum(f32[4096]{0} %param_5.544, f32[4096]{0} %add.711), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.283 = f32[4096]{0} select(pred[4096]{0} %param_3.832, f32[4096]{0} %maximum.85, f32[4096]{0} %broadcast.1688), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %broadcast.1687 = f32[4096]{0} broadcast(f32[] %constant_500), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.282 = f32[4096]{0} select(pred[4096]{0} %param_3.832, f32[4096]{0} %param_4.605, f32[4096]{0} %broadcast.1687), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.248 = f32[4096]{0} subtract(f32[4096]{0} %select.283, f32[4096]{0} %select.282), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1239 = f32[] constant(0.00452488707)
  %broadcast.1686 = f32[4096]{0} broadcast(f32[] %constant_1239), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.508 = f32[4096]{0} multiply(f32[4096]{0} %subtract.248, f32[4096]{0} %broadcast.1686), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.1685 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.508), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.507 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.35, f32[4096,221]{1,0} %broadcast.1685), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %exponential.20 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %multiply.507), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %subtract.247 = f32[4096,221]{1,0} subtract(f32[4096,221]{1,0} %broadcast.1689, f32[4096,221]{1,0} %exponential.20), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %param_1.779 = f32[4096,221]{1,0} parameter(1)
  %multiply.263 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %subtract.247, f32[4096,221]{1,0} %param_1.779), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=280}
  %slice.305 = f32[4096,220]{1,0} slice(f32[4096,221]{1,0} %multiply.263), slice={[0:4096], [1:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=283}
  %pad.107 = f32[4096,221]{1,0} pad(f32[4096,220]{1,0} %slice.305, f32[] %constant_500), padding=0_0x0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0, 1), inserted_window_dims=(), scatter_dims_to_operand_dims=(1,)) indices_are_sorted=True unique_indices=True mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=286}
  %add.364 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %add.365, f32[4096,221]{1,0} %pad.107), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=286}
  %param_0.62 = f32[4096,221]{1,0} parameter(0)
  ROOT %multiply.262 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %add.364, f32[4096,221]{1,0} %param_0.62), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%region_16.1504 (Arg_0.1505: f32[], Arg_1.1506: f32[]) -> f32[] {
  %Arg_0.1505 = f32[] parameter(0)
  %Arg_1.1506 = f32[] parameter(1)
  ROOT %add.1507 = f32[] add(f32[] %Arg_0.1505, f32[] %Arg_1.1506), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.6 (param_0.63: f32[4096,221], param_1.88: s32[4096,33,2], param_2.926: f32[4096], param_3.827: f32[4096,221], param_4.601: f32[4096]) -> f32[4096,221] {
  %param_0.63 = f32[4096,221]{1,0} parameter(0)
  %param_1.88 = s32[4096,33,2]{1,0,2} parameter(1)
  %param_2.926 = f32[4096]{0} parameter(2)
  %constant_501 = f32[] constant(1)
  %broadcast.815 = f32[4096]{0} broadcast(f32[] %constant_501), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_4.601 = f32[4096]{0} parameter(4)
  %constant_1727 = f32[] constant(1e-08)
  %broadcast.2197 = f32[4096]{0} broadcast(f32[] %constant_1727), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.950 = f32[4096]{0} add(f32[4096]{0} %param_4.601, f32[4096]{0} %broadcast.2197), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  %multiply.266 = f32[4096]{0} multiply(f32[4096]{0} %add.950, f32[4096]{0} %add.950), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %divide.28 = f32[4096]{0} divide(f32[4096]{0} %broadcast.815, f32[4096]{0} %multiply.266), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %multiply.265 = f32[4096]{0} multiply(f32[4096]{0} %param_2.926, f32[4096]{0} %divide.28), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %broadcast.2193 = f32[4096,1]{1,0} broadcast(f32[] %constant_501), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_3.827 = f32[4096,221]{1,0} parameter(3)
  %slice.654 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_3.827), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(1,), start_index_map=(1,)) slice_sizes=(4096, 1) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=482}
  %subtract.378 = f32[4096,1]{1,0} subtract(f32[4096,1]{1,0} %broadcast.2193, f32[4096,1]{1,0} %slice.654), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %broadcast.2192 = f32[4096,1]{1,0} broadcast(f32[] %constant_1727), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.946 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %subtract.378, f32[4096,1]{1,0} %broadcast.2192), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %bitcast.550 = f32[4096]{0} bitcast(f32[4096,1]{1,0} %add.946), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %multiply.264 = f32[4096]{0} multiply(f32[4096]{0} %multiply.265, f32[4096]{0} %bitcast.550), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %negate.19 = f32[4096]{0} negate(f32[4096]{0} %multiply.264), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %broadcast.814 = f32[4096,33]{1,0} broadcast(f32[4096]{0} %negate.19), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/broadcast_in_dim[shape=(4096, 33) broadcast_dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  ROOT %scatter.48 = f32[4096,221]{1,0} scatter(f32[4096,221]{1,0} %param_0.63, s32[4096,33,2]{1,0,2} %param_1.88, f32[4096,33]{1,0} %broadcast.814), update_window_dims={}, inserted_window_dims={0,1}, scatter_dims_to_operand_dims={0,1}, index_vector_dim=2, to_apply=%region_16.1504, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(), inserted_window_dims=(0, 1), scatter_dims_to_operand_dims=(0, 1)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=472}
}

%region_18.1512 (Arg_0.1513: f32[], Arg_1.1514: f32[]) -> f32[] {
  %Arg_0.1513 = f32[] parameter(0)
  %Arg_1.1514 = f32[] parameter(1)
  ROOT %add.1515 = f32[] add(f32[] %Arg_0.1513, f32[] %Arg_1.1514), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
}

%fused_computation.24 (param_0.384: f32[4096,221,3], param_1.921: f32[4096,3], param_2.1052: f32[4096], param_3.933: f32[4096,221]) -> f32[4096,221] {
  %param_0.384 = f32[4096,221,3]{1,0,2} parameter(0)
  %param_1.921 = f32[4096,3]{1,0} parameter(1)
  %constant_1834 = f32[] constant(1)
  %broadcast.2288 = f32[4096,1]{1,0} broadcast(f32[] %constant_1834), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_3.933 = f32[4096,221]{1,0} parameter(3)
  %slice.712 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_3.933), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(1,), start_index_map=(1,)) slice_sizes=(4096, 1) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=482}
  %subtract.396 = f32[4096,1]{1,0} subtract(f32[4096,1]{1,0} %broadcast.2288, f32[4096,1]{1,0} %slice.712), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %constant_1833 = f32[] constant(1e-08)
  %broadcast.2287 = f32[4096,1]{1,0} broadcast(f32[] %constant_1833), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.1028 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %subtract.396, f32[4096,1]{1,0} %broadcast.2287), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %bitcast.560 = f32[4096]{0} bitcast(f32[4096,1]{1,0} %add.1028), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_2.1052 = f32[4096]{0} parameter(2)
  %broadcast.2286 = f32[4096]{0} broadcast(f32[] %constant_1833), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.1027 = f32[4096]{0} add(f32[4096]{0} %param_2.1052, f32[4096]{0} %broadcast.2286), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  %divide.64 = f32[4096]{0} divide(f32[4096]{0} %bitcast.560, f32[4096]{0} %add.1027), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %constant_1740 = f32[] constant(0.000162760422)
  %broadcast.2205 = f32[4096]{0} broadcast(f32[] %constant_1740), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.581 = f32[4096]{0} multiply(f32[4096]{0} %divide.64, f32[4096]{0} %broadcast.2205), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %broadcast.2204 = f32[4096,3]{1,0} broadcast(f32[4096]{0} %multiply.581), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.580 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %param_1.921, f32[4096,3]{1,0} %broadcast.2204), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %broadcast.2203 = f32[4096,221,3]{1,0,2} broadcast(f32[4096,3]{1,0} %multiply.580), dimensions={0,2}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/broadcast_in_dim[shape=(4096, 221, 3) broadcast_dimensions=(0, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.267 = f32[4096,221,3]{1,0,2} multiply(f32[4096,221,3]{1,0,2} %param_0.384, f32[4096,221,3]{1,0,2} %broadcast.2203), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %bitcast.78 = f32[3,4096,221]{2,1,0} bitcast(f32[4096,221,3]{1,0,2} %multiply.267), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %constant_504 = f32[] constant(0)
  ROOT %reduce.38 = f32[4096,221]{1,0} reduce(f32[3,4096,221]{2,1,0} %bitcast.78, f32[] %constant_504), dimensions={0}, to_apply=%region_18.1512, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
}

%region_14.1496 (Arg_0.1497: f32[], Arg_1.1498: f32[]) -> f32[] {
  %Arg_0.1497 = f32[] parameter(0)
  %Arg_1.1498 = f32[] parameter(1)
  ROOT %add.1499 = f32[] add(f32[] %Arg_0.1497, f32[] %Arg_1.1498), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
}

%fused_computation.25 (param_0.716: f32[3,4096], param_1.925: f32[4096,3]) -> f32[4096] {
  %param_0.716 = f32[3,4096]{0,1} parameter(0)
  %bitcast.137 = f32[4096,3]{1,0} bitcast(f32[3,4096]{0,1} %param_0.716), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %param_1.925 = f32[4096,3]{1,0} parameter(1)
  %constant_1748 = f32[] constant(0.000162760422)
  %broadcast.2213 = f32[4096,3]{1,0} broadcast(f32[] %constant_1748), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %multiply.587 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %param_1.925, f32[4096,3]{1,0} %broadcast.2213), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %multiply.268 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %bitcast.137, f32[4096,3]{1,0} %multiply.587), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %constant_502 = f32[] constant(0)
  ROOT %reduce.39 = f32[4096]{0} reduce(f32[4096,3]{1,0} %multiply.268, f32[] %constant_502), dimensions={1}, to_apply=%region_14.1496, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
}

%fused_computation.27 (param_0.69: f32[3,48,128,128], param_1.969: f32[], param_2.995: f32[], param_3.878: f32[], param_4.641: f32[3,48,128,128], param_5.587: f32[3,48,128,128], param_6.526: f32[3,48,128,128], param_7.461: f32[3,48,128,128]) -> (f32[3,48,128,128], f32[3,48,128,128], f32[3,48,128,128]) {
  %param_0.69 = f32[3,48,128,128]{3,2,1,0} parameter(0)
  %param_5.587 = f32[3,48,128,128]{3,2,1,0} parameter(5)
  %param_6.526 = f32[3,48,128,128]{3,2,1,0} parameter(6)
  %add.495.clone.1 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %param_5.587, f32[3,48,128,128]{3,2,1,0} %param_6.526), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_510_clone_1 = f32[] constant(0.0100000007)
  %broadcast.828.clone.1 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_510_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.277.clone.1 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %add.495.clone.1, f32[3,48,128,128]{3,2,1,0} %broadcast.828.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_7.461 = f32[3,48,128,128]{3,2,1,0} parameter(7)
  %constant_511_clone_1 = f32[] constant(0.9)
  %broadcast.826.clone.1 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_511_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.276.clone.1 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %param_7.461, f32[3,48,128,128]{3,2,1,0} %broadcast.826.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.370.clone.1 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %multiply.277.clone.1, f32[3,48,128,128]{3,2,1,0} %multiply.276.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.995 = f32[] parameter(2)
  %broadcast.821 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %param_2.995), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_509_clone_1 = f32[] constant(0.1)
  %broadcast.824.clone.1 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_509_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.275.clone.1 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %add.495.clone.1, f32[3,48,128,128]{3,2,1,0} %broadcast.824.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.274.clone.1 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %multiply.275.clone.1, f32[3,48,128,128]{3,2,1,0} %multiply.275.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_507_clone_1 = f32[] constant(0.01)
  %broadcast.823.clone.1 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_507_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.273.clone.1 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %multiply.274.clone.1, f32[3,48,128,128]{3,2,1,0} %broadcast.823.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.641 = f32[3,48,128,128]{3,2,1,0} parameter(4)
  %constant_508_clone_1 = f32[] constant(0.99)
  %broadcast.822.clone.1 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_508_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.272.clone.1 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %param_4.641, f32[3,48,128,128]{3,2,1,0} %broadcast.822.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.369.clone.1 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %multiply.273.clone.1, f32[3,48,128,128]{3,2,1,0} %multiply.272.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.969 = f32[] parameter(1)
  %broadcast.819 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %param_1.969), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.30 = f32[3,48,128,128]{3,2,1,0} divide(f32[3,48,128,128]{3,2,1,0} %add.369.clone.1, f32[3,48,128,128]{3,2,1,0} %broadcast.819), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.2 = f32[3,48,128,128]{3,2,1,0} sqrt(f32[3,48,128,128]{3,2,1,0} %divide.30), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_506 = f32[] constant(1e-08)
  %broadcast.818 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_506), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.368 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %sqrt.2, f32[3,48,128,128]{3,2,1,0} %broadcast.818), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.271 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %broadcast.821, f32[3,48,128,128]{3,2,1,0} %add.368)
  %divide.29 = f32[3,48,128,128]{3,2,1,0} divide(f32[3,48,128,128]{3,2,1,0} %add.370.clone.1, f32[3,48,128,128]{3,2,1,0} %multiply.271), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.878 = f32[] parameter(3)
  %constant_505 = f32[] constant(-0.02)
  %multiply.399 = f32[] multiply(f32[] %param_3.878, f32[] %constant_505), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.817 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %multiply.399), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.270 = f32[3,48,128,128]{3,2,1,0} multiply(f32[3,48,128,128]{3,2,1,0} %divide.29, f32[3,48,128,128]{3,2,1,0} %broadcast.817), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.366 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %param_0.69, f32[3,48,128,128]{3,2,1,0} %multiply.270), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.52 = (f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}) tuple(f32[3,48,128,128]{3,2,1,0} %add.366, f32[3,48,128,128]{3,2,1,0} %add.369.clone.1, f32[3,48,128,128]{3,2,1,0} %add.370.clone.1)
}

%region_35.1749 (Arg_0.1750: f32[], Arg_1.1751: f32[]) -> f32[] {
  %Arg_0.1750 = f32[] parameter(0)
  %Arg_1.1751 = f32[] parameter(1)
  ROOT %add.1752 = f32[] add(f32[] %Arg_0.1750, f32[] %Arg_1.1751), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.7 (param_0.691: f32[3,4096,48,33], param_1.901: f32[3,4096,3,33]) -> f32[3,48,128,128] {
  %constant_512 = f32[] constant(0)
  %broadcast.969 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_512), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/broadcast_in_dim[shape=(3, 48, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.183 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1571 = s32[] constant(0)
  %broadcast.2007 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1571), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.901 = f32[3,4096,3,33]{3,1,0,2} parameter(1)
  %slice.564 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.901), slice={[0:3], [0:4096], [1:2], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 4096, 1, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.446 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.564), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.332 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.446), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.194 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.332), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1570 = s32[] constant(127)
  %broadcast.2006 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1570), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.183 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.2007, s32[3,4096,33]{2,1,0} %convert.194, s32[3,4096,33]{2,1,0} %broadcast.2006), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.311 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.183, s32[3,4096,33]{2,1,0} %broadcast.2007), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1569 = s32[] constant(128)
  %broadcast.2005 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1569), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.863 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.183, s32[3,4096,33]{2,1,0} %broadcast.2005), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.357 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.311, s32[3,4096,33]{2,1,0} %add.863, s32[3,4096,33]{2,1,0} %clamp.183), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.445 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.357), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.561 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.901), slice={[0:3], [0:4096], [2:3], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 1, 0) limit_indices=(3, 4096, 2, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.444 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.561), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.331 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.444), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.193 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.331), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.182 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.2007, s32[3,4096,33]{2,1,0} %convert.193, s32[3,4096,33]{2,1,0} %broadcast.2006), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.310 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.182, s32[3,4096,33]{2,1,0} %broadcast.2007), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.862 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.182, s32[3,4096,33]{2,1,0} %broadcast.2005), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.356 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.310, s32[3,4096,33]{2,1,0} %add.862, s32[3,4096,33]{2,1,0} %clamp.182), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.443 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.356), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.235 = s32[3,4096,33,3]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.183, s32[3,4096,33,1]{2,1,0,3} %bitcast.445, s32[3,4096,33,1]{2,1,0,3} %bitcast.443), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.24 = s32[3,4096,33,3]{3,2,1,0} copy(s32[3,4096,33,3]{2,1,0,3} %concatenate.235), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1548 = f32[] constant(1)
  %broadcast.1981 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1548), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.322 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.446, f32[3,4096,33]{2,1,0} %floor.332), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.321 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1981, f32[3,4096,33]{2,1,0} %subtract.322), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.320 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.444, f32[3,4096,33]{2,1,0} %floor.331), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.319 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1981, f32[3,4096,33]{2,1,0} %subtract.320), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.540 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.321, f32[3,4096,33]{2,1,0} %subtract.319), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1980 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.540), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.691 = f32[3,4096,48,33]{3,2,1,0} parameter(0)
  %multiply.278 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.1980, f32[3,4096,48,33]{3,2,1,0} %param_0.691), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.49 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.969, s32[3,4096,33,3]{3,2,1,0} %copy.24, f32[3,4096,48,33]{3,2,1,0} %multiply.278), update_window_dims={2}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_35.1749, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_34.1745 (Arg_0.1746: f32[], Arg_1.1747: f32[]) -> f32[] {
  %Arg_0.1746 = f32[] parameter(0)
  %Arg_1.1747 = f32[] parameter(1)
  ROOT %add.1748 = f32[] add(f32[] %Arg_0.1746, f32[] %Arg_1.1747), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.8 (param_0.688: f32[3,4096,48,33], param_1.897: f32[3,48,128,128], param_2.896: f32[3,48,128,128], param_3.773: f32[3,4096,3,33]) -> f32[3,48,128,128] {
  %param_1.897 = f32[3,48,128,128]{3,2,1,0} parameter(1)
  %param_2.896 = f32[3,48,128,128]{3,2,1,0} parameter(2)
  %add.371 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %param_1.897, f32[3,48,128,128]{3,2,1,0} %param_2.896), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.179 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1539 = s32[] constant(0)
  %broadcast.1973 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1539), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_3.773 = f32[3,4096,3,33]{3,1,0,2} parameter(3)
  %slice.540 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_3.773), slice={[0:3], [0:4096], [1:2], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 4096, 1, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.422 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.540), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.316 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.422), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.186 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.316), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1537 = s32[] constant(127)
  %broadcast.1972 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1537), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.175 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1973, s32[3,4096,33]{2,1,0} %convert.186, s32[3,4096,33]{2,1,0} %broadcast.1972), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.303 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.175, s32[3,4096,33]{2,1,0} %broadcast.1973), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1536 = s32[] constant(128)
  %broadcast.1970 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1536), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.855 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.175, s32[3,4096,33]{2,1,0} %broadcast.1970), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.349 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.303, s32[3,4096,33]{2,1,0} %add.855, s32[3,4096,33]{2,1,0} %clamp.175), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.421 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.349), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.539 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_3.773), slice={[0:3], [0:4096], [2:3], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 1, 0) limit_indices=(3, 4096, 2, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.420 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.539), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.315 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.420), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.185 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.315), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1538 = s32[] constant(1)
  %broadcast.1968 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1538), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.854 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.185, s32[3,4096,33]{2,1,0} %broadcast.1968), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.174 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1973, s32[3,4096,33]{2,1,0} %add.854, s32[3,4096,33]{2,1,0} %broadcast.1972), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.302 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.174, s32[3,4096,33]{2,1,0} %broadcast.1973), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.853 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.174, s32[3,4096,33]{2,1,0} %broadcast.1970), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.348 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.302, s32[3,4096,33]{2,1,0} %add.853, s32[3,4096,33]{2,1,0} %clamp.174), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.419 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.348), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.231 = s32[3,4096,33,3]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.179, s32[3,4096,33,1]{3,2,1,0} %bitcast.421, s32[3,4096,33,1]{3,2,1,0} %bitcast.419), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1497 = f32[] constant(1)
  %broadcast.1924 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1497), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.306 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.422, f32[3,4096,33]{2,1,0} %floor.316), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.305 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1924, f32[3,4096,33]{2,1,0} %subtract.306), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.304 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.420, f32[3,4096,33]{2,1,0} %floor.315), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.536 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.305, f32[3,4096,33]{2,1,0} %subtract.304), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1923 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.536), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.688 = f32[3,4096,48,33]{3,2,1,0} parameter(0)
  %multiply.279 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.1923, f32[3,4096,48,33]{3,2,1,0} %param_0.688), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.50 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %add.371, s32[3,4096,33,3]{3,2,1,0} %concatenate.231, f32[3,4096,48,33]{3,2,1,0} %multiply.279), update_window_dims={2}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_34.1745, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_33.1741 (Arg_0.1742: f32[], Arg_1.1743: f32[]) -> f32[] {
  %Arg_0.1742 = f32[] parameter(0)
  %Arg_1.1743 = f32[] parameter(1)
  ROOT %add.1744 = f32[] add(f32[] %Arg_0.1742, f32[] %Arg_1.1743), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.9 (param_0.673: f32[3,4096,48,33], param_1.885: f32[3,4096,3,33]) -> f32[3,48,128,128] {
  %constant_513 = f32[] constant(0)
  %broadcast.970 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_513), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/broadcast_in_dim[shape=(3, 48, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.175 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1487 = s32[] constant(0)
  %broadcast.1911 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1487), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.885 = f32[3,4096,3,33]{3,1,0,2} parameter(1)
  %slice.520 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.885), slice={[0:3], [0:4096], [1:2], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 4096, 1, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.394 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.520), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.296 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.394), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.174 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.296), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1490 = s32[] constant(1)
  %broadcast.1910 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1490), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.837 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.174, s32[3,4096,33]{2,1,0} %broadcast.1910), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1486 = s32[] constant(127)
  %broadcast.1909 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1486), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.163 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1911, s32[3,4096,33]{2,1,0} %add.837, s32[3,4096,33]{2,1,0} %broadcast.1909), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.290 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.163, s32[3,4096,33]{2,1,0} %broadcast.1911), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1485 = s32[] constant(128)
  %broadcast.1908 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1485), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.836 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.163, s32[3,4096,33]{2,1,0} %broadcast.1908), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.337 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.290, s32[3,4096,33]{2,1,0} %add.836, s32[3,4096,33]{2,1,0} %clamp.163), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.393 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.337), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.519 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.885), slice={[0:3], [0:4096], [2:3], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 1, 0) limit_indices=(3, 4096, 2, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.392 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.519), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.295 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.392), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.173 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.295), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.162 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1911, s32[3,4096,33]{2,1,0} %convert.173, s32[3,4096,33]{2,1,0} %broadcast.1909), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.289 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.162, s32[3,4096,33]{2,1,0} %broadcast.1911), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.835 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.162, s32[3,4096,33]{2,1,0} %broadcast.1908), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.336 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.289, s32[3,4096,33]{2,1,0} %add.835, s32[3,4096,33]{2,1,0} %clamp.162), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.391 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.336), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.227 = s32[3,4096,33,3]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.175, s32[3,4096,33,1]{2,1,0,3} %bitcast.393, s32[3,4096,33,1]{2,1,0,3} %bitcast.391), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.25 = s32[3,4096,33,3]{3,2,1,0} copy(s32[3,4096,33,3]{2,1,0,3} %concatenate.227), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.294 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.394, f32[3,4096,33]{2,1,0} %floor.296), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1447 = f32[] constant(1)
  %broadcast.1871 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1447), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.293 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.392, f32[3,4096,33]{2,1,0} %floor.295), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.292 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1871, f32[3,4096,33]{2,1,0} %subtract.293), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.532 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.294, f32[3,4096,33]{2,1,0} %subtract.292), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1870 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.532), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.673 = f32[3,4096,48,33]{3,2,1,0} parameter(0)
  %multiply.280 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.1870, f32[3,4096,48,33]{3,2,1,0} %param_0.673), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.51 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.970, s32[3,4096,33,3]{3,2,1,0} %copy.25, f32[3,4096,48,33]{3,2,1,0} %multiply.280), update_window_dims={2}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_33.1741, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_32.1737 (Arg_0.1738: f32[], Arg_1.1739: f32[]) -> f32[] {
  %Arg_0.1738 = f32[] parameter(0)
  %Arg_1.1739 = f32[] parameter(1)
  ROOT %add.1740 = f32[] add(f32[] %Arg_0.1738, f32[] %Arg_1.1739), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.10 (param_0.670: f32[3,4096,48,33], param_1.871: f32[3,4096,3,33]) -> f32[3,48,128,128] {
  %constant_514 = f32[] constant(0)
  %broadcast.971 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_514), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/broadcast_in_dim[shape=(3, 48, 128, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.171 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1437 = s32[] constant(0)
  %broadcast.1863 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1437), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.871 = f32[3,4096,3,33]{3,1,0,2} parameter(1)
  %slice.500 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.871), slice={[0:3], [0:4096], [1:2], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 4096, 1, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.362 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.500), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.275 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.362), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.162 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.275), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1436 = s32[] constant(1)
  %broadcast.1862 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1436), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.819 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.162, s32[3,4096,33]{2,1,0} %broadcast.1862), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1435 = s32[] constant(127)
  %broadcast.1861 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1435), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.151 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1863, s32[3,4096,33]{2,1,0} %add.819, s32[3,4096,33]{2,1,0} %broadcast.1861), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.277 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.151, s32[3,4096,33]{2,1,0} %broadcast.1863), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1434 = s32[] constant(128)
  %broadcast.1860 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1434), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.818 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.151, s32[3,4096,33]{2,1,0} %broadcast.1860), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.325 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.277, s32[3,4096,33]{2,1,0} %add.818, s32[3,4096,33]{2,1,0} %clamp.151), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.361 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.325), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.499 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.871), slice={[0:3], [0:4096], [2:3], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 1, 0) limit_indices=(3, 4096, 2, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.360 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.499), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.274 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.360), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.161 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.274), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.817 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.161, s32[3,4096,33]{2,1,0} %broadcast.1862), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.150 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1863, s32[3,4096,33]{2,1,0} %add.817, s32[3,4096,33]{2,1,0} %broadcast.1861), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.276 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.150, s32[3,4096,33]{2,1,0} %broadcast.1863), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.816 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.150, s32[3,4096,33]{2,1,0} %broadcast.1860), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.324 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.276, s32[3,4096,33]{2,1,0} %add.816, s32[3,4096,33]{2,1,0} %clamp.150), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.359 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.324), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.223 = s32[3,4096,33,3]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.171, s32[3,4096,33,1]{3,2,1,0} %bitcast.361, s32[3,4096,33,1]{3,2,1,0} %bitcast.359), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.282 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.362, f32[3,4096,33]{2,1,0} %floor.275), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.279 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.360, f32[3,4096,33]{2,1,0} %floor.274), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.528 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.282, f32[3,4096,33]{2,1,0} %subtract.279), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1794 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.528), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.670 = f32[3,4096,48,33]{3,2,1,0} parameter(0)
  %multiply.281 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.1794, f32[3,4096,48,33]{3,2,1,0} %param_0.670), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.52 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.971, s32[3,4096,33,3]{3,2,1,0} %concatenate.223, f32[3,4096,48,33]{3,2,1,0} %multiply.281), update_window_dims={2}, inserted_window_dims={0,2,3}, scatter_dims_to_operand_dims={0,2,3}, index_vector_dim=3, to_apply=%region_32.1737, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%fused_computation.30 (param_0.714: f32[3,48,128], param_1.924: f32[3,4096,3,33], param_2.940: f32[135168,144]) -> f32[3,4096,48,33] {
  %constant_1340 = f32[] constant(1)
  %broadcast.1761 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1340), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.924 = f32[3,4096,3,33]{3,1,0,2} parameter(1)
  %slice.442 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.924), slice={[0:3], [0:4096], [0:1], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1, 2, 3), collapsed_slice_dims=(), start_index_map=(2,)) slice_sizes=(3, 4096, 1, 33) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=165}
  %bitcast.290 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.442), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.243 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.290), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.266 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.290, f32[3,4096,33]{2,1,0} %floor.243), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.265 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1761, f32[3,4096,33]{2,1,0} %subtract.266), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1760 = f32[3,4096,48,33]{2,1,0,3} broadcast(f32[3,4096,33]{2,1,0} %subtract.265), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.714 = f32[3,48,128]{2,1,0} parameter(0)
  %iota.163 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1336 = s32[] constant(0)
  %broadcast.1759 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1336), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.142 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.243), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1334 = s32[] constant(127)
  %broadcast.1758 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1334), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.131 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1759, s32[3,4096,33]{2,1,0} %convert.142, s32[3,4096,33]{2,1,0} %broadcast.1758), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.257 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.131, s32[3,4096,33]{2,1,0} %broadcast.1759), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1333 = s32[] constant(128)
  %broadcast.1757 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1333), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.773 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.131, s32[3,4096,33]{2,1,0} %broadcast.1757), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.305 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.257, s32[3,4096,33]{2,1,0} %add.773, s32[3,4096,33]{2,1,0} %clamp.131), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.288 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.305), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.214 = s32[3,4096,33,2]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.163, s32[3,4096,33,1]{2,1,0,3} %bitcast.288), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.107 = f32[3,4096,48,33]{2,1,0,3} gather(f32[3,48,128]{2,1,0} %param_0.714, s32[3,4096,33,2]{2,1,0,3} %concatenate.214), offset_dims={2}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,48,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 48, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.520 = f32[3,4096,48,33]{2,1,0,3} multiply(f32[3,4096,48,33]{2,1,0,3} %broadcast.1760, f32[3,4096,48,33]{2,1,0,3} %gather.107), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.37 = f32[3,4096,48,33]{3,2,1,0} copy(f32[3,4096,48,33]{2,1,0,3} %multiply.520), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1756 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %subtract.266), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.162 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1335 = s32[] constant(1)
  %broadcast.1754 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1335), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.772 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.142, s32[3,4096,33]{2,1,0} %broadcast.1754), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.130 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1759, s32[3,4096,33]{2,1,0} %add.772, s32[3,4096,33]{2,1,0} %broadcast.1758), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.256 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.130, s32[3,4096,33]{2,1,0} %broadcast.1759), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.771 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.130, s32[3,4096,33]{2,1,0} %broadcast.1757), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.304 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.256, s32[3,4096,33]{2,1,0} %add.771, s32[3,4096,33]{2,1,0} %clamp.130), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.285 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.304), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.213 = s32[3,4096,33,2]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.162, s32[3,4096,33,1]{3,2,1,0} %bitcast.285), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.106 = f32[3,4096,48,33]{3,2,1,0} gather(f32[3,48,128]{2,1,0} %param_0.714, s32[3,4096,33,2]{3,2,1,0} %concatenate.213), offset_dims={2}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,48,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 48, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.519 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.1756, f32[3,4096,48,33]{3,2,1,0} %gather.106), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.770 = f32[3,4096,48,33]{3,2,1,0} add(f32[3,4096,48,33]{3,2,1,0} %copy.37, f32[3,4096,48,33]{3,2,1,0} %multiply.519), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.940 = f32[135168,144]{1,0} parameter(2)
  %bitcast.554 = f32[4096,33,144]{2,1,0} bitcast(f32[135168,144]{1,0} %param_2.940), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reshape[new_sizes=(4096, 33, 144) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
  %transpose.47 = f32[144,4096,33]{0,2,1} transpose(f32[4096,33,144]{2,1,0} %bitcast.554), dimensions={2,0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/transpose[permutation=(2, 0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
  %reshape.592 = f32[3,48,4096,33]{1,2,0,3} reshape(f32[144,4096,33]{0,2,1} %transpose.47), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 48, 4096, 33) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %transpose.46 = f32[3,4096,48,33]{2,1,0,3} transpose(f32[3,48,4096,33]{1,2,0,3} %reshape.592), dimensions={0,2,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/transpose[permutation=(0, 2, 1, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=90}
  %copy.26 = f32[3,4096,48,33]{3,2,1,0} copy(f32[3,4096,48,33]{2,1,0,3} %transpose.46), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/transpose[permutation=(0, 2, 1, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=90}
  ROOT %multiply.282 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %add.770, f32[3,4096,48,33]{3,2,1,0} %copy.26), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
}

%fused_computation.31 (param_0.80: f32[3,48,128], param_1.970: f32[], param_2.997: f32[], param_3.880: f32[], param_4.643: f32[3,48,128], param_5.590: f32[3,48,128], param_6.533: f32[3,48,128], param_7.470: f32[3,48,128]) -> (f32[3,48,128], f32[3,48,128], f32[3,48,128]) {
  %param_0.80 = f32[3,48,128]{2,1,0} parameter(0)
  %param_5.590 = f32[3,48,128]{2,1,0} parameter(5)
  %param_6.533 = f32[3,48,128]{2,1,0} parameter(6)
  %add.497.clone.1 = f32[3,48,128]{2,1,0} add(f32[3,48,128]{2,1,0} %param_5.590, f32[3,48,128]{2,1,0} %param_6.533), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/add_any" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_520_clone_1 = f32[] constant(0.0100000007)
  %broadcast.837.clone.1 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_520_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.290.clone.1 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %add.497.clone.1, f32[3,48,128]{2,1,0} %broadcast.837.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_7.470 = f32[3,48,128]{2,1,0} parameter(7)
  %constant_521_clone_1 = f32[] constant(0.9)
  %broadcast.836.clone.1 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_521_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.289.clone.1 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %param_7.470, f32[3,48,128]{2,1,0} %broadcast.836.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.376.clone.1 = f32[3,48,128]{2,1,0} add(f32[3,48,128]{2,1,0} %multiply.290.clone.1, f32[3,48,128]{2,1,0} %multiply.289.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.997 = f32[] parameter(2)
  %broadcast.832 = f32[3,48,128]{2,1,0} broadcast(f32[] %param_2.997), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_519_clone_1 = f32[] constant(0.1)
  %broadcast.835.clone.1 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_519_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.288.clone.1 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %add.497.clone.1, f32[3,48,128]{2,1,0} %broadcast.835.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.287.clone.1 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %multiply.288.clone.1, f32[3,48,128]{2,1,0} %multiply.288.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_517_clone_1 = f32[] constant(0.01)
  %broadcast.834.clone.1 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_517_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.286.clone.1 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %multiply.287.clone.1, f32[3,48,128]{2,1,0} %broadcast.834.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.643 = f32[3,48,128]{2,1,0} parameter(4)
  %constant_518_clone_1 = f32[] constant(0.99)
  %broadcast.833.clone.1 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_518_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.285.clone.1 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %param_4.643, f32[3,48,128]{2,1,0} %broadcast.833.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.375.clone.1 = f32[3,48,128]{2,1,0} add(f32[3,48,128]{2,1,0} %multiply.286.clone.1, f32[3,48,128]{2,1,0} %multiply.285.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.970 = f32[] parameter(1)
  %broadcast.831 = f32[3,48,128]{2,1,0} broadcast(f32[] %param_1.970), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.32 = f32[3,48,128]{2,1,0} divide(f32[3,48,128]{2,1,0} %add.375.clone.1, f32[3,48,128]{2,1,0} %broadcast.831), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.3 = f32[3,48,128]{2,1,0} sqrt(f32[3,48,128]{2,1,0} %divide.32), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_516 = f32[] constant(1e-08)
  %broadcast.830 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_516), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.373 = f32[3,48,128]{2,1,0} add(f32[3,48,128]{2,1,0} %sqrt.3, f32[3,48,128]{2,1,0} %broadcast.830), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.284 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %broadcast.832, f32[3,48,128]{2,1,0} %add.373)
  %divide.31 = f32[3,48,128]{2,1,0} divide(f32[3,48,128]{2,1,0} %add.376.clone.1, f32[3,48,128]{2,1,0} %multiply.284), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.880 = f32[] parameter(3)
  %constant_515 = f32[] constant(-0.02)
  %multiply.400 = f32[] multiply(f32[] %param_3.880, f32[] %constant_515), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.829 = f32[3,48,128]{2,1,0} broadcast(f32[] %multiply.400), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.283 = f32[3,48,128]{2,1,0} multiply(f32[3,48,128]{2,1,0} %divide.31, f32[3,48,128]{2,1,0} %broadcast.829), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.372 = f32[3,48,128]{2,1,0} add(f32[3,48,128]{2,1,0} %param_0.80, f32[3,48,128]{2,1,0} %multiply.283), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.54 = (f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}) tuple(f32[3,48,128]{2,1,0} %add.372, f32[3,48,128]{2,1,0} %add.375.clone.1, f32[3,48,128]{2,1,0} %add.376.clone.1)
}

%region_37.1792 (Arg_0.1793: f32[], Arg_1.1794: f32[]) -> f32[] {
  %Arg_0.1793 = f32[] parameter(0)
  %Arg_1.1794 = f32[] parameter(1)
  ROOT %add.1795 = f32[] add(f32[] %Arg_0.1793, f32[] %Arg_1.1794), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.11 (param_0.625: f32[3,4096,48,33], param_1.831: f32[3,4096,3,33]) -> f32[3,48,128] {
  %constant_522 = f32[] constant(0)
  %broadcast.972 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_522), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/broadcast_in_dim[shape=(3, 48, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.159 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1323 = s32[] constant(0)
  %broadcast.1741 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1323), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.831 = f32[3,4096,3,33]{3,1,0,2} parameter(1)
  %slice.434 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.831), slice={[0:3], [0:4096], [0:1], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1, 2, 3), collapsed_slice_dims=(), start_index_map=(2,)) slice_sizes=(3, 4096, 1, 33) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=165}
  %bitcast.278 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.434), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.235 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.278), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.138 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.235), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1322 = s32[] constant(127)
  %broadcast.1740 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1322), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.127 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1741, s32[3,4096,33]{2,1,0} %convert.138, s32[3,4096,33]{2,1,0} %broadcast.1740), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.253 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.127, s32[3,4096,33]{2,1,0} %broadcast.1741), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1321 = s32[] constant(128)
  %broadcast.1739 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1321), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.765 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.127, s32[3,4096,33]{2,1,0} %broadcast.1739), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.300 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.253, s32[3,4096,33]{2,1,0} %add.765, s32[3,4096,33]{2,1,0} %clamp.127), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.277 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.300), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.210 = s32[3,4096,33,2]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.159, s32[3,4096,33,1]{2,1,0,3} %bitcast.277), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1309 = f32[] constant(1)
  %broadcast.1729 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1309), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.260 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.278, f32[3,4096,33]{2,1,0} %floor.235), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.259 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1729, f32[3,4096,33]{2,1,0} %subtract.260), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1728 = f32[3,4096,48,33]{2,1,0,3} broadcast(f32[3,4096,33]{2,1,0} %subtract.259), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.625 = f32[3,4096,48,33]{2,1,0,3} parameter(0)
  %multiply.291 = f32[3,4096,48,33]{2,1,0,3} multiply(f32[3,4096,48,33]{2,1,0,3} %broadcast.1728, f32[3,4096,48,33]{2,1,0,3} %param_0.625), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.53 = f32[3,48,128]{2,1,0} scatter(f32[3,48,128]{2,1,0} %broadcast.972, s32[3,4096,33,2]{2,1,0,3} %concatenate.210, f32[3,4096,48,33]{2,1,0,3} %multiply.291), update_window_dims={2}, inserted_window_dims={0,2}, scatter_dims_to_operand_dims={0,2}, index_vector_dim=3, to_apply=%region_37.1792, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%region_36.1788 (Arg_0.1789: f32[], Arg_1.1790: f32[]) -> f32[] {
  %Arg_0.1789 = f32[] parameter(0)
  %Arg_1.1790 = f32[] parameter(1)
  ROOT %add.1791 = f32[] add(f32[] %Arg_0.1789, f32[] %Arg_1.1790), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.12 (param_0.620: f32[3,4096,48,33], param_1.827: f32[3,4096,3,33]) -> f32[3,48,128] {
  %constant_523 = f32[] constant(0)
  %broadcast.973 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_523), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/broadcast_in_dim[shape=(3, 48, 128) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.155 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1303 = s32[] constant(0)
  %broadcast.1721 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1303), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.827 = f32[3,4096,3,33]{3,1,0,2} parameter(1)
  %slice.426 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.827), slice={[0:3], [0:4096], [0:1], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1, 2, 3), collapsed_slice_dims=(), start_index_map=(2,)) slice_sizes=(3, 4096, 1, 33) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=165}
  %bitcast.266 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.426), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.227 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.266), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.134 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.227), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1302 = s32[] constant(1)
  %broadcast.1720 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1302), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.761 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.134, s32[3,4096,33]{2,1,0} %broadcast.1720), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1301 = s32[] constant(127)
  %broadcast.1719 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1301), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.123 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1721, s32[3,4096,33]{2,1,0} %add.761, s32[3,4096,33]{2,1,0} %broadcast.1719), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.249 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.123, s32[3,4096,33]{2,1,0} %broadcast.1721), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1300 = s32[] constant(128)
  %broadcast.1718 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1300), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.760 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.123, s32[3,4096,33]{2,1,0} %broadcast.1718), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.296 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.249, s32[3,4096,33]{2,1,0} %add.760, s32[3,4096,33]{2,1,0} %clamp.123), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.265 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.296), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.206 = s32[3,4096,33,2]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.155, s32[3,4096,33,1]{3,2,1,0} %bitcast.265), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.78 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.266, f32[3,4096,33]{2,1,0} %floor.227), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.974 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %subtract.78), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.620 = f32[3,4096,48,33]{2,1,0,3} parameter(0)
  %copy.27 = f32[3,4096,48,33]{3,2,1,0} copy(f32[3,4096,48,33]{2,1,0,3} %param_0.620), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %multiply.292 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.974, f32[3,4096,48,33]{3,2,1,0} %copy.27), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %scatter.54 = f32[3,48,128]{2,1,0} scatter(f32[3,48,128]{2,1,0} %broadcast.973, s32[3,4096,33,2]{3,2,1,0} %concatenate.206, f32[3,4096,48,33]{3,2,1,0} %multiply.292), update_window_dims={2}, inserted_window_dims={0,2}, scatter_dims_to_operand_dims={0,2}, index_vector_dim=3, to_apply=%region_36.1788, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%fused_computation.34 (param_0.712: f32[3,48,128,128], param_1.923: f32[3,4096,3,33], param_2.936: f32[135168,144]) -> f32[3,4096,48,33] {
  %param_2.936 = f32[135168,144]{1,0} parameter(2)
  %bitcast.552 = f32[4096,33,144]{2,1,0} bitcast(f32[135168,144]{1,0} %param_2.936), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reshape[new_sizes=(4096, 33, 144) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
  %transpose.43 = f32[144,4096,33]{0,2,1} transpose(f32[4096,33,144]{2,1,0} %bitcast.552), dimensions={2,0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/transpose[permutation=(2, 0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
  %reshape.590 = f32[3,48,4096,33]{1,2,0,3} reshape(f32[144,4096,33]{0,2,1} %transpose.43), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/reshape[new_sizes=(3, 48, 4096, 33) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %transpose.42 = f32[3,4096,48,33]{2,1,0,3} transpose(f32[3,48,4096,33]{1,2,0,3} %reshape.590), dimensions={0,2,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/transpose[permutation=(0, 2, 1, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=90}
  %constant_1622 = f32[] constant(1)
  %broadcast.2077 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1622), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.923 = f32[3,4096,3,33]{3,1,0,2} parameter(1)
  %slice.616 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.923), slice={[0:3], [0:4096], [1:2], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 4096, 1, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.494 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.616), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.365 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.494), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.346 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.494, f32[3,4096,33]{2,1,0} %floor.365), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.345 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.2077, f32[3,4096,33]{2,1,0} %subtract.346), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.615 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.923), slice={[0:3], [0:4096], [2:3], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 1, 0) limit_indices=(3, 4096, 2, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.493 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.615), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.364 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.493), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.344 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.493, f32[3,4096,33]{2,1,0} %floor.364), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.343 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.2077, f32[3,4096,33]{2,1,0} %subtract.344), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.557 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.345, f32[3,4096,33]{2,1,0} %subtract.343), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2076 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.557), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.712 = f32[3,48,128,128]{3,2,1,0} parameter(0)
  %iota.191 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1610 = s32[] constant(0)
  %broadcast.2075 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1610), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.210 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.365), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1608 = s32[] constant(127)
  %broadcast.2074 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1608), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.199 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.2075, s32[3,4096,33]{2,1,0} %convert.210, s32[3,4096,33]{2,1,0} %broadcast.2074), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.327 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.199, s32[3,4096,33]{2,1,0} %broadcast.2075), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1607 = s32[] constant(128)
  %broadcast.2073 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1607), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.898 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.199, s32[3,4096,33]{2,1,0} %broadcast.2073), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.373 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.327, s32[3,4096,33]{2,1,0} %add.898, s32[3,4096,33]{2,1,0} %clamp.199), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.491 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.373), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.209 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.364), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.198 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.2075, s32[3,4096,33]{2,1,0} %convert.209, s32[3,4096,33]{2,1,0} %broadcast.2074), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.326 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.198, s32[3,4096,33]{2,1,0} %broadcast.2075), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.897 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.198, s32[3,4096,33]{2,1,0} %broadcast.2073), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.372 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.326, s32[3,4096,33]{2,1,0} %add.897, s32[3,4096,33]{2,1,0} %clamp.198), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.489 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.372), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.243 = s32[3,4096,33,3]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.191, s32[3,4096,33,1]{2,1,0,3} %bitcast.491, s32[3,4096,33,1]{2,1,0,3} %bitcast.489), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.119 = f32[3,4096,48,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_0.712, s32[3,4096,33,3]{2,1,0,3} %concatenate.243), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.556 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.2076, f32[3,4096,48,33]{3,2,1,0} %gather.119), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.555 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.345, f32[3,4096,33]{2,1,0} %subtract.344), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2068 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.555), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.190 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.485 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.373), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1609 = s32[] constant(1)
  %broadcast.2063 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1609), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.895 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.209, s32[3,4096,33]{2,1,0} %broadcast.2063), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.196 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.2075, s32[3,4096,33]{2,1,0} %add.895, s32[3,4096,33]{2,1,0} %broadcast.2074), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.324 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.196, s32[3,4096,33]{2,1,0} %broadcast.2075), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.894 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.196, s32[3,4096,33]{2,1,0} %broadcast.2073), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.370 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.324, s32[3,4096,33]{2,1,0} %add.894, s32[3,4096,33]{2,1,0} %clamp.196), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.483 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.370), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.242 = s32[3,4096,33,3]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.190, s32[3,4096,33,1]{3,2,1,0} %bitcast.485, s32[3,4096,33,1]{3,2,1,0} %bitcast.483), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.118 = f32[3,4096,48,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_0.712, s32[3,4096,33,3]{3,2,1,0} %concatenate.242), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.554 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.2068, f32[3,4096,48,33]{3,2,1,0} %gather.118), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.893 = f32[3,4096,48,33]{3,2,1,0} add(f32[3,4096,48,33]{3,2,1,0} %multiply.556, f32[3,4096,48,33]{3,2,1,0} %multiply.554), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.553 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.346, f32[3,4096,33]{2,1,0} %subtract.343), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2059 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.553), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.189 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.892 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.210, s32[3,4096,33]{2,1,0} %broadcast.2063), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.195 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.2075, s32[3,4096,33]{2,1,0} %add.892, s32[3,4096,33]{2,1,0} %broadcast.2074), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.323 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.195, s32[3,4096,33]{2,1,0} %broadcast.2075), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.890 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.195, s32[3,4096,33]{2,1,0} %broadcast.2073), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.369 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.323, s32[3,4096,33]{2,1,0} %add.890, s32[3,4096,33]{2,1,0} %clamp.195), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.479 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.369), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.241 = s32[3,4096,33,3]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.189, s32[3,4096,33,1]{2,1,0,3} %bitcast.479, s32[3,4096,33,1]{2,1,0,3} %bitcast.489), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.117 = f32[3,4096,48,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_0.712, s32[3,4096,33,3]{2,1,0,3} %concatenate.241), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.552 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.2059, f32[3,4096,48,33]{3,2,1,0} %gather.117), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.888 = f32[3,4096,48,33]{3,2,1,0} add(f32[3,4096,48,33]{3,2,1,0} %add.893, f32[3,4096,48,33]{3,2,1,0} %multiply.552), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.551 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.346, f32[3,4096,33]{2,1,0} %subtract.344), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2051 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.551), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.188 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.473 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.369), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.240 = s32[3,4096,33,3]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.188, s32[3,4096,33,1]{3,2,1,0} %bitcast.473, s32[3,4096,33,1]{3,2,1,0} %bitcast.483), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.116 = f32[3,4096,48,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_0.712, s32[3,4096,33,3]{3,2,1,0} %concatenate.240), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.550 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.2051, f32[3,4096,48,33]{3,2,1,0} %gather.116), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.882 = f32[3,4096,48,33]{3,2,1,0} add(f32[3,4096,48,33]{3,2,1,0} %add.888, f32[3,4096,48,33]{3,2,1,0} %multiply.550), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.28 = f32[3,4096,48,33]{2,1,0,3} copy(f32[3,4096,48,33]{3,2,1,0} %add.882), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %multiply.293 = f32[3,4096,48,33]{2,1,0,3} multiply(f32[3,4096,48,33]{2,1,0,3} %transpose.42, f32[3,4096,48,33]{2,1,0,3} %copy.28), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
}

%fused_computation.36 (param_0.92: f32[128,3], param_1.971: f32[], param_2.999: f32[], param_3.882: f32[], param_4.645: f32[128,3], param_5.593: f32[128,3], param_6.540: f32[128,3]) -> (f32[128,3], f32[128,3], f32[128,3]) {
  %param_0.92 = f32[128,3]{1,0} parameter(0)
  %param_5.593 = f32[128,3]{1,0} parameter(5)
  %constant_529_clone_1 = f32[] constant(0.0100000007)
  %broadcast.849.clone.1 = f32[128,3]{1,0} broadcast(f32[] %constant_529_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.302.clone.1 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %param_5.593, f32[128,3]{1,0} %broadcast.849.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_6.540 = f32[128,3]{1,0} parameter(6)
  %constant_530_clone_1 = f32[] constant(0.9)
  %broadcast.848.clone.1 = f32[128,3]{1,0} broadcast(f32[] %constant_530_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.301.clone.1 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %param_6.540, f32[128,3]{1,0} %broadcast.848.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.382.clone.1 = f32[128,3]{1,0} add(f32[128,3]{1,0} %multiply.302.clone.1, f32[128,3]{1,0} %multiply.301.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.999 = f32[] parameter(2)
  %broadcast.842 = f32[128,3]{1,0} broadcast(f32[] %param_2.999), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_528_clone_1 = f32[] constant(0.1)
  %broadcast.847.clone.1 = f32[128,3]{1,0} broadcast(f32[] %constant_528_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.300.clone.1 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %param_5.593, f32[128,3]{1,0} %broadcast.847.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.299.clone.1 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %multiply.300.clone.1, f32[128,3]{1,0} %multiply.300.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_526_clone_1 = f32[] constant(0.01)
  %broadcast.846.clone.1 = f32[128,3]{1,0} broadcast(f32[] %constant_526_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.297.clone.1 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %multiply.299.clone.1, f32[128,3]{1,0} %broadcast.846.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.645 = f32[128,3]{1,0} parameter(4)
  %constant_527_clone_1 = f32[] constant(0.99)
  %broadcast.844.clone.1 = f32[128,3]{1,0} broadcast(f32[] %constant_527_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.296.clone.1 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %param_4.645, f32[128,3]{1,0} %broadcast.844.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.380.clone.1 = f32[128,3]{1,0} add(f32[128,3]{1,0} %multiply.297.clone.1, f32[128,3]{1,0} %multiply.296.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.971 = f32[] parameter(1)
  %broadcast.841 = f32[128,3]{1,0} broadcast(f32[] %param_1.971), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.34 = f32[128,3]{1,0} divide(f32[128,3]{1,0} %add.380.clone.1, f32[128,3]{1,0} %broadcast.841), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.4 = f32[128,3]{1,0} sqrt(f32[128,3]{1,0} %divide.34), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_524 = f32[] constant(1e-08)
  %broadcast.840 = f32[128,3]{1,0} broadcast(f32[] %constant_524), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.379 = f32[128,3]{1,0} add(f32[128,3]{1,0} %sqrt.4, f32[128,3]{1,0} %broadcast.840), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.295 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %broadcast.842, f32[128,3]{1,0} %add.379)
  %divide.33 = f32[128,3]{1,0} divide(f32[128,3]{1,0} %add.382.clone.1, f32[128,3]{1,0} %multiply.295), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.882 = f32[] parameter(3)
  %constant_525 = f32[] constant(-0.001)
  %multiply.401 = f32[] multiply(f32[] %param_3.882, f32[] %constant_525), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.839 = f32[128,3]{1,0} broadcast(f32[] %multiply.401), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.294 = f32[128,3]{1,0} multiply(f32[128,3]{1,0} %divide.33, f32[128,3]{1,0} %broadcast.839), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.377 = f32[128,3]{1,0} add(f32[128,3]{1,0} %param_0.92, f32[128,3]{1,0} %multiply.294), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.56 = (f32[128,3]{1,0}, f32[128,3]{1,0}, f32[128,3]{1,0}) tuple(f32[128,3]{1,0} %add.377, f32[128,3]{1,0} %add.380.clone.1, f32[128,3]{1,0} %add.382.clone.1)
}

%fused_computation.39 (param_0.97: f32[3], param_1.972: f32[], param_2.1001: f32[], param_3.884: f32[], param_4.647: f32[3], param_5.596: f32[3], param_6.547: f32[3]) -> (f32[3], f32[3], f32[3]) {
  %param_0.97 = f32[3]{0} parameter(0)
  %param_5.596 = f32[3]{0} parameter(5)
  %constant_539_clone_1 = f32[] constant(0.0100000007)
  %broadcast.859.clone.1 = f32[3]{0} broadcast(f32[] %constant_539_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.310.clone.1 = f32[3]{0} multiply(f32[3]{0} %param_5.596, f32[3]{0} %broadcast.859.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_6.547 = f32[3]{0} parameter(6)
  %constant_540_clone_1 = f32[] constant(0.9)
  %broadcast.858.clone.1 = f32[3]{0} broadcast(f32[] %constant_540_clone_1), dimensions={}
  %multiply.309.clone.1 = f32[3]{0} multiply(f32[3]{0} %param_6.547, f32[3]{0} %broadcast.858.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.386.clone.1 = f32[3]{0} add(f32[3]{0} %multiply.310.clone.1, f32[3]{0} %multiply.309.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.1001 = f32[] parameter(2)
  %broadcast.853 = f32[3]{0} broadcast(f32[] %param_2.1001), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_538_clone_1 = f32[] constant(0.1)
  %broadcast.857.clone.1 = f32[3]{0} broadcast(f32[] %constant_538_clone_1), dimensions={}
  %multiply.308.clone.1 = f32[3]{0} multiply(f32[3]{0} %param_5.596, f32[3]{0} %broadcast.857.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.307.clone.1 = f32[3]{0} multiply(f32[3]{0} %multiply.308.clone.1, f32[3]{0} %multiply.308.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_536_clone_1 = f32[] constant(0.01)
  %broadcast.855.clone.1 = f32[3]{0} broadcast(f32[] %constant_536_clone_1), dimensions={}
  %multiply.306.clone.1 = f32[3]{0} multiply(f32[3]{0} %multiply.307.clone.1, f32[3]{0} %broadcast.855.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.647 = f32[3]{0} parameter(4)
  %constant_537_clone_1 = f32[] constant(0.99)
  %broadcast.854.clone.1 = f32[3]{0} broadcast(f32[] %constant_537_clone_1), dimensions={}
  %multiply.305.clone.1 = f32[3]{0} multiply(f32[3]{0} %param_4.647, f32[3]{0} %broadcast.854.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.385.clone.1 = f32[3]{0} add(f32[3]{0} %multiply.306.clone.1, f32[3]{0} %multiply.305.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.972 = f32[] parameter(1)
  %broadcast.852 = f32[3]{0} broadcast(f32[] %param_1.972), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.36 = f32[3]{0} divide(f32[3]{0} %add.385.clone.1, f32[3]{0} %broadcast.852), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.5 = f32[3]{0} sqrt(f32[3]{0} %divide.36), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_531 = f32[] constant(1e-08)
  %broadcast.851 = f32[3]{0} broadcast(f32[] %constant_531), dimensions={}
  %add.384 = f32[3]{0} add(f32[3]{0} %sqrt.5, f32[3]{0} %broadcast.851), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.304 = f32[3]{0} multiply(f32[3]{0} %broadcast.853, f32[3]{0} %add.384)
  %divide.35 = f32[3]{0} divide(f32[3]{0} %add.386.clone.1, f32[3]{0} %multiply.304), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.884 = f32[] parameter(3)
  %constant_533 = f32[] constant(-0.001)
  %multiply.402 = f32[] multiply(f32[] %param_3.884, f32[] %constant_533), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.850 = f32[3]{0} broadcast(f32[] %multiply.402), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.303 = f32[3]{0} multiply(f32[3]{0} %divide.35, f32[3]{0} %broadcast.850), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.383 = f32[3]{0} add(f32[3]{0} %param_0.97, f32[3]{0} %multiply.303), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.58 = (f32[3]{0}, f32[3]{0}, f32[3]{0}) tuple(f32[3]{0} %add.383, f32[3]{0} %add.385.clone.1, f32[3]{0} %add.386.clone.1)
}

%input_fused_computation_reduce.1 (param_0.390: f32[135168,3]) -> f32[3,368] {
  %param_0.390 = f32[135168,3]{0,1} parameter(0)
  %bitcast.81 = f32[3,135168]{1,0} bitcast(f32[135168,3]{0,1} %param_0.390), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(sigmoid)))/jit(jit_transpose(jvp(sigmoid)))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %constant_541 = f32[] constant(0)
  %pad.110 = f32[3,135424]{1,0} pad(f32[3,135168]{1,0} %bitcast.81, f32[] %constant_541), padding=0_0x0_256
  %bitcast.80 = f32[3,368,368]{2,1,0} bitcast(f32[3,135424]{1,0} %pad.110)
  ROOT %reduce.40 = f32[3,368]{1,0} reduce(f32[3,368,368]{2,1,0} %bitcast.80, f32[] %constant_541), dimensions={2}, to_apply=%region_27.1717
}

%fused_computation.42 (param_0.104: f32[128,128], param_1.973: f32[], param_2.1003: f32[], param_3.886: f32[], param_4.649: f32[128,128], param_5.599: f32[128,128], param_6.554: f32[128,128]) -> (f32[128,128], f32[128,128], f32[128,128]) {
  %param_0.104 = f32[128,128]{1,0} parameter(0)
  %param_5.599 = f32[128,128]{1,0} parameter(5)
  %constant_547_clone_1 = f32[] constant(0.0100000007)
  %broadcast.870.clone.1 = f32[128,128]{1,0} broadcast(f32[] %constant_547_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.319.clone.1 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %param_5.599, f32[128,128]{1,0} %broadcast.870.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_6.554 = f32[128,128]{1,0} parameter(6)
  %constant_548_clone_1 = f32[] constant(0.9)
  %broadcast.869.clone.1 = f32[128,128]{1,0} broadcast(f32[] %constant_548_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.318.clone.1 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %param_6.554, f32[128,128]{1,0} %broadcast.869.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.390.clone.1 = f32[128,128]{1,0} add(f32[128,128]{1,0} %multiply.319.clone.1, f32[128,128]{1,0} %multiply.318.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.1003 = f32[] parameter(2)
  %broadcast.865 = f32[128,128]{1,0} broadcast(f32[] %param_2.1003), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_546_clone_1 = f32[] constant(0.1)
  %broadcast.868.clone.1 = f32[128,128]{1,0} broadcast(f32[] %constant_546_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.317.clone.1 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %param_5.599, f32[128,128]{1,0} %broadcast.868.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.316.clone.1 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %multiply.317.clone.1, f32[128,128]{1,0} %multiply.317.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_544_clone_1 = f32[] constant(0.01)
  %broadcast.867.clone.1 = f32[128,128]{1,0} broadcast(f32[] %constant_544_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.315.clone.1 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %multiply.316.clone.1, f32[128,128]{1,0} %broadcast.867.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.649 = f32[128,128]{1,0} parameter(4)
  %constant_545_clone_1 = f32[] constant(0.99)
  %broadcast.866.clone.1 = f32[128,128]{1,0} broadcast(f32[] %constant_545_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.314.clone.1 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %param_4.649, f32[128,128]{1,0} %broadcast.866.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.389.clone.1 = f32[128,128]{1,0} add(f32[128,128]{1,0} %multiply.315.clone.1, f32[128,128]{1,0} %multiply.314.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.973 = f32[] parameter(1)
  %broadcast.864 = f32[128,128]{1,0} broadcast(f32[] %param_1.973), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.38 = f32[128,128]{1,0} divide(f32[128,128]{1,0} %add.389.clone.1, f32[128,128]{1,0} %broadcast.864), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.6 = f32[128,128]{1,0} sqrt(f32[128,128]{1,0} %divide.38), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_542 = f32[] constant(1e-08)
  %broadcast.862 = f32[128,128]{1,0} broadcast(f32[] %constant_542), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.388 = f32[128,128]{1,0} add(f32[128,128]{1,0} %sqrt.6, f32[128,128]{1,0} %broadcast.862), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.313 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %broadcast.865, f32[128,128]{1,0} %add.388)
  %divide.37 = f32[128,128]{1,0} divide(f32[128,128]{1,0} %add.390.clone.1, f32[128,128]{1,0} %multiply.313), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.886 = f32[] parameter(3)
  %constant_543 = f32[] constant(-0.001)
  %multiply.403 = f32[] multiply(f32[] %param_3.886, f32[] %constant_543), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.860 = f32[128,128]{1,0} broadcast(f32[] %multiply.403), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.312 = f32[128,128]{1,0} multiply(f32[128,128]{1,0} %divide.37, f32[128,128]{1,0} %broadcast.860), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.387 = f32[128,128]{1,0} add(f32[128,128]{1,0} %param_0.104, f32[128,128]{1,0} %multiply.312), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.60 = (f32[128,128]{1,0}, f32[128,128]{1,0}, f32[128,128]{1,0}) tuple(f32[128,128]{1,0} %add.387, f32[128,128]{1,0} %add.389.clone.1, f32[128,128]{1,0} %add.390.clone.1)
}

%fused_computation.45 (param_0.109: f32[128], param_1.974: f32[], param_2.1005: f32[], param_3.888: f32[], param_4.651: f32[128], param_5.602: f32[128], param_6.561: f32[128], param_7.534: f32[128], param_8.362: f32[128], param_9.254: f32[128], param_10.146: f32[128]) -> (f32[128], f32[128], f32[128], f32[128], f32[128], /*index=5*/f32[128]) {
  %param_0.109 = f32[128]{0} parameter(0)
  %param_5.602 = f32[128]{0} parameter(5)
  %constant_555_clone_1 = f32[] constant(0.0100000007)
  %broadcast.982.clone.1 = f32[128]{0} broadcast(f32[] %constant_555_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.327.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_5.602, f32[128]{0} %broadcast.982.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_6.561 = f32[128]{0} parameter(6)
  %constant_556_clone_1 = f32[] constant(0.9)
  %broadcast.983.clone.1 = f32[128]{0} broadcast(f32[] %constant_556_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.326.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_6.561, f32[128]{0} %broadcast.983.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.394.clone.1 = f32[128]{0} add(f32[128]{0} %multiply.327.clone.1, f32[128]{0} %multiply.326.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.1005 = f32[] parameter(2)
  %broadcast.977 = f32[128]{0} broadcast(f32[] %param_2.1005), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_554_clone_1 = f32[] constant(0.1)
  %broadcast.979.clone.1 = f32[128]{0} broadcast(f32[] %constant_554_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.325.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_5.602, f32[128]{0} %broadcast.979.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.324.clone.1 = f32[128]{0} multiply(f32[128]{0} %multiply.325.clone.1, f32[128]{0} %multiply.325.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_551_clone_1 = f32[] constant(0.01)
  %broadcast.980.clone.1 = f32[128]{0} broadcast(f32[] %constant_551_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.323.clone.1 = f32[128]{0} multiply(f32[128]{0} %multiply.324.clone.1, f32[128]{0} %broadcast.980.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.651 = f32[128]{0} parameter(4)
  %constant_553_clone_1 = f32[] constant(0.99)
  %broadcast.981.clone.1 = f32[128]{0} broadcast(f32[] %constant_553_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.322.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_4.651, f32[128]{0} %broadcast.981.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.393.clone.1 = f32[128]{0} add(f32[128]{0} %multiply.323.clone.1, f32[128]{0} %multiply.322.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.974 = f32[] parameter(1)
  %broadcast.975 = f32[128]{0} broadcast(f32[] %param_1.974), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.40 = f32[128]{0} divide(f32[128]{0} %add.393.clone.1, f32[128]{0} %broadcast.975), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.7 = f32[128]{0} sqrt(f32[128]{0} %divide.40), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_549 = f32[] constant(1e-08)
  %broadcast.976 = f32[128]{0} broadcast(f32[] %constant_549), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.392 = f32[128]{0} add(f32[128]{0} %sqrt.7, f32[128]{0} %broadcast.976), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.321 = f32[128]{0} multiply(f32[128]{0} %broadcast.977, f32[128]{0} %add.392)
  %divide.39 = f32[128]{0} divide(f32[128]{0} %add.394.clone.1, f32[128]{0} %multiply.321), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.888 = f32[] parameter(3)
  %constant_550 = f32[] constant(-0.001)
  %multiply.404 = f32[] multiply(f32[] %param_3.888, f32[] %constant_550), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.978 = f32[128]{0} broadcast(f32[] %multiply.404), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.320 = f32[128]{0} multiply(f32[128]{0} %divide.39, f32[128]{0} %broadcast.978), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.391 = f32[128]{0} add(f32[128]{0} %param_0.109, f32[128]{0} %multiply.320), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %param_7.534 = f32[128]{0} parameter(7)
  %param_9.254 = f32[128]{0} parameter(9)
  %multiply.343.clone.1.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_9.254, f32[128]{0} %broadcast.982.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_10.146 = f32[128]{0} parameter(10)
  %multiply.342.clone.1.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_10.146, f32[128]{0} %broadcast.983.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.404.clone.1.clone.1 = f32[128]{0} add(f32[128]{0} %multiply.343.clone.1.clone.1, f32[128]{0} %multiply.342.clone.1.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.341.clone.1.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_9.254, f32[128]{0} %broadcast.979.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.340.clone.1.clone.1 = f32[128]{0} multiply(f32[128]{0} %multiply.341.clone.1.clone.1, f32[128]{0} %multiply.341.clone.1.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.339.clone.1.clone.1 = f32[128]{0} multiply(f32[128]{0} %multiply.340.clone.1.clone.1, f32[128]{0} %broadcast.980.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_8.362 = f32[128]{0} parameter(8)
  %multiply.338.clone.1.clone.1 = f32[128]{0} multiply(f32[128]{0} %param_8.362, f32[128]{0} %broadcast.981.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.403.clone.1.clone.1 = f32[128]{0} add(f32[128]{0} %multiply.339.clone.1.clone.1, f32[128]{0} %multiply.338.clone.1.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %divide.44.clone.1 = f32[128]{0} divide(f32[128]{0} %add.403.clone.1.clone.1, f32[128]{0} %broadcast.975), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.9.clone.1 = f32[128]{0} sqrt(f32[128]{0} %divide.44.clone.1), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.402.clone.1 = f32[128]{0} add(f32[128]{0} %sqrt.9.clone.1, f32[128]{0} %broadcast.976), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.337.clone.1 = f32[128]{0} multiply(f32[128]{0} %broadcast.977, f32[128]{0} %add.402.clone.1)
  %divide.43.clone.1 = f32[128]{0} divide(f32[128]{0} %add.404.clone.1.clone.1, f32[128]{0} %multiply.337.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.336.clone.1 = f32[128]{0} multiply(f32[128]{0} %divide.43.clone.1, f32[128]{0} %broadcast.978), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.401.clone.1 = f32[128]{0} add(f32[128]{0} %param_7.534, f32[128]{0} %multiply.336.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.67 = (f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) tuple(f32[128]{0} %add.391, f32[128]{0} %add.393.clone.1, f32[128]{0} %add.394.clone.1, f32[128]{0} %add.401.clone.1, f32[128]{0} %add.403.clone.1.clone.1, /*index=5*/f32[128]{0} %add.404.clone.1.clone.1)
}

%region_28.1721 (Arg_0.1722: f32[], Arg_1.1723: f32[]) -> f32[] {
  %Arg_0.1722 = f32[] parameter(0)
  %Arg_1.1723 = f32[] parameter(1)
  ROOT %add.1724 = f32[] add(f32[] %Arg_0.1722, f32[] %Arg_1.1723), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
}

%region_29.1725 (Arg_0.1726: f32[], Arg_1.1727: f32[]) -> f32[] {
  %Arg_0.1726 = f32[] parameter(0)
  %Arg_1.1727 = f32[] parameter(1)
  ROOT %add.1728 = f32[] add(f32[] %Arg_0.1726, f32[] %Arg_1.1727), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
}

%input_fused_computation_reduce.2 (param_0.391: f32[135168,128], param_1.1009: f32[135168,128]) -> (f32[368,128], f32[368,128]) {
  %param_0.391 = f32[135168,128]{1,0} parameter(0)
  %constant_557 = f32[] constant(0)
  %pad.111 = f32[135424,128]{1,0} pad(f32[135168,128]{1,0} %param_0.391, f32[] %constant_557), padding=0_256x0_0
  %bitcast.82 = f32[368,368,128]{2,1,0} bitcast(f32[135424,128]{1,0} %pad.111)
  %reduce.41 = f32[368,128]{1,0} reduce(f32[368,368,128]{2,1,0} %bitcast.82, f32[] %constant_557), dimensions={0}, to_apply=%region_28.1721
  %param_1.1009 = f32[135168,128]{1,0} parameter(1)
  %pad.112.clone.1 = f32[135424,128]{1,0} pad(f32[135168,128]{1,0} %param_1.1009, f32[] %constant_557), padding=0_256x0_0
  %bitcast.83.clone.1 = f32[368,368,128]{2,1,0} bitcast(f32[135424,128]{1,0} %pad.112.clone.1)
  %reduce.42.clone.1 = f32[368,128]{1,0} reduce(f32[368,368,128]{2,1,0} %bitcast.83.clone.1, f32[] %constant_557), dimensions={0}, to_apply=%region_29.1725
  ROOT %tuple.90 = (f32[368,128]{1,0}, f32[368,128]{1,0}) tuple(f32[368,128]{1,0} %reduce.41, f32[368,128]{1,0} %reduce.42.clone.1)
}

%fused_computation.48 (param_0.116: f32[150,128], param_1.975: f32[], param_2.1007: f32[], param_3.890: f32[], param_4.653: f32[150,128], param_5.605: f32[150,128], param_6.568: f32[150,128]) -> (f32[150,128], f32[150,128], f32[150,128]) {
  %param_0.116 = f32[150,128]{1,0} parameter(0)
  %param_5.605 = f32[150,128]{1,0} parameter(5)
  %constant_563_clone_1 = f32[] constant(0.0100000007)
  %broadcast.883.clone.1 = f32[150,128]{1,0} broadcast(f32[] %constant_563_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.335.clone.1 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %param_5.605, f32[150,128]{1,0} %broadcast.883.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_6.568 = f32[150,128]{1,0} parameter(6)
  %constant_564_clone_1 = f32[] constant(0.9)
  %broadcast.881.clone.1 = f32[150,128]{1,0} broadcast(f32[] %constant_564_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.334.clone.1 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %param_6.568, f32[150,128]{1,0} %broadcast.881.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.400.clone.1 = f32[150,128]{1,0} add(f32[150,128]{1,0} %multiply.335.clone.1, f32[150,128]{1,0} %multiply.334.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.1007 = f32[] parameter(2)
  %broadcast.874 = f32[150,128]{1,0} broadcast(f32[] %param_2.1007), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_562_clone_1 = f32[] constant(0.1)
  %broadcast.879.clone.1 = f32[150,128]{1,0} broadcast(f32[] %constant_562_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.333.clone.1 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %param_5.605, f32[150,128]{1,0} %broadcast.879.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.332.clone.1 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %multiply.333.clone.1, f32[150,128]{1,0} %multiply.333.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_560_clone_1 = f32[] constant(0.01)
  %broadcast.877.clone.1 = f32[150,128]{1,0} broadcast(f32[] %constant_560_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.331.clone.1 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %multiply.332.clone.1, f32[150,128]{1,0} %broadcast.877.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.653 = f32[150,128]{1,0} parameter(4)
  %constant_561_clone_1 = f32[] constant(0.99)
  %broadcast.875.clone.1 = f32[150,128]{1,0} broadcast(f32[] %constant_561_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.330.clone.1 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %param_4.653, f32[150,128]{1,0} %broadcast.875.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.398.clone.1 = f32[150,128]{1,0} add(f32[150,128]{1,0} %multiply.331.clone.1, f32[150,128]{1,0} %multiply.330.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.975 = f32[] parameter(1)
  %broadcast.873 = f32[150,128]{1,0} broadcast(f32[] %param_1.975), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.42 = f32[150,128]{1,0} divide(f32[150,128]{1,0} %add.398.clone.1, f32[150,128]{1,0} %broadcast.873), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.8 = f32[150,128]{1,0} sqrt(f32[150,128]{1,0} %divide.42), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_558 = f32[] constant(1e-08)
  %broadcast.872 = f32[150,128]{1,0} broadcast(f32[] %constant_558), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.397 = f32[150,128]{1,0} add(f32[150,128]{1,0} %sqrt.8, f32[150,128]{1,0} %broadcast.872), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.329 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %broadcast.874, f32[150,128]{1,0} %add.397)
  %divide.41 = f32[150,128]{1,0} divide(f32[150,128]{1,0} %add.400.clone.1, f32[150,128]{1,0} %multiply.329), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.890 = f32[] parameter(3)
  %constant_559 = f32[] constant(-0.001)
  %multiply.405 = f32[] multiply(f32[] %param_3.890, f32[] %constant_559), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.871 = f32[150,128]{1,0} broadcast(f32[] %multiply.405), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.328 = f32[150,128]{1,0} multiply(f32[150,128]{1,0} %divide.41, f32[150,128]{1,0} %broadcast.871), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.395 = f32[150,128]{1,0} add(f32[150,128]{1,0} %param_0.116, f32[150,128]{1,0} %multiply.328), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.64 = (f32[150,128]{1,0}, f32[150,128]{1,0}, f32[150,128]{1,0}) tuple(f32[150,128]{1,0} %add.395, f32[150,128]{1,0} %add.398.clone.1, f32[150,128]{1,0} %add.400.clone.1)
}

%fused_computation.54 (param_0.128: f32[144,27], param_1.986: f32[], param_2.1021: f32[], param_3.898: f32[], param_4.661: f32[144,27], param_5.612: f32[144,27], param_6.582: f32[144,27]) -> (f32[144,27], f32[144,27], f32[144,27]) {
  %param_0.128 = f32[144,27]{1,0} parameter(0)
  %param_5.612 = f32[144,27]{1,0} parameter(5)
  %constant_583_clone_1 = f32[] constant(0.0100000007)
  %broadcast.893.clone.1 = f32[144,27]{1,0} broadcast(f32[] %constant_583_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.352.clone.1 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %param_5.612, f32[144,27]{1,0} %broadcast.893.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_6.582 = f32[144,27]{1,0} parameter(6)
  %constant_584_clone_1 = f32[] constant(0.9)
  %broadcast.892.clone.1 = f32[144,27]{1,0} broadcast(f32[] %constant_584_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.351.clone.1 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %param_6.582, f32[144,27]{1,0} %broadcast.892.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.409.clone.1 = f32[144,27]{1,0} add(f32[144,27]{1,0} %multiply.352.clone.1, f32[144,27]{1,0} %multiply.351.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_2.1021 = f32[] parameter(2)
  %broadcast.888 = f32[144,27]{1,0} broadcast(f32[] %param_2.1021), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %constant_582_clone_1 = f32[] constant(0.1)
  %broadcast.891.clone.1 = f32[144,27]{1,0} broadcast(f32[] %constant_582_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.350.clone.1 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %param_5.612, f32[144,27]{1,0} %broadcast.891.clone.1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=325}
  %multiply.349.clone.1 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %multiply.350.clone.1, f32[144,27]{1,0} %multiply.350.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %constant_580_clone_1 = f32[] constant(0.01)
  %broadcast.890.clone.1 = f32[144,27]{1,0} broadcast(f32[] %constant_580_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.348.clone.1 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %multiply.349.clone.1, f32[144,27]{1,0} %broadcast.890.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_4.661 = f32[144,27]{1,0} parameter(4)
  %constant_581_clone_1 = f32[] constant(0.99)
  %broadcast.889.clone.1 = f32[144,27]{1,0} broadcast(f32[] %constant_581_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %multiply.347.clone.1 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %param_4.661, f32[144,27]{1,0} %broadcast.889.clone.1), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %add.407.clone.1 = f32[144,27]{1,0} add(f32[144,27]{1,0} %multiply.348.clone.1, f32[144,27]{1,0} %multiply.347.clone.1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=81}
  %param_1.986 = f32[] parameter(1)
  %broadcast.887 = f32[144,27]{1,0} broadcast(f32[] %param_1.986), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %divide.46 = f32[144,27]{1,0} divide(f32[144,27]{1,0} %add.407.clone.1, f32[144,27]{1,0} %broadcast.887), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=87}
  %sqrt.10 = f32[144,27]{1,0} sqrt(f32[144,27]{1,0} %divide.46), metadata={op_name="jit(training_step)/jit(main)/sqrt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %constant_573 = f32[] constant(1e-08)
  %broadcast.886 = f32[144,27]{1,0} broadcast(f32[] %constant_573), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %add.406 = f32[144,27]{1,0} add(f32[144,27]{1,0} %sqrt.10, f32[144,27]{1,0} %broadcast.886), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %multiply.345 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %broadcast.888, f32[144,27]{1,0} %add.406)
  %divide.45 = f32[144,27]{1,0} divide(f32[144,27]{1,0} %add.409.clone.1, f32[144,27]{1,0} %multiply.345), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=302}
  %param_3.898 = f32[] parameter(3)
  %constant_574 = f32[] constant(-0.001)
  %multiply.408 = f32[] multiply(f32[] %param_3.898, f32[] %constant_574), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %broadcast.885 = f32[144,27]{1,0} broadcast(f32[] %multiply.408), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %multiply.344 = f32[144,27]{1,0} multiply(f32[144,27]{1,0} %divide.45, f32[144,27]{1,0} %broadcast.885), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=347}
  %add.405 = f32[144,27]{1,0} add(f32[144,27]{1,0} %param_0.128, f32[144,27]{1,0} %multiply.344), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  ROOT %tuple.73 = (f32[144,27]{1,0}, f32[144,27]{1,0}, f32[144,27]{1,0}) tuple(f32[144,27]{1,0} %add.405, f32[144,27]{1,0} %add.407.clone.1, f32[144,27]{1,0} %add.409.clone.1)
}

%fused_computation.56 (param_0.750: s32[]) -> (f32[], f32[], s32[]) {
  %constant_578 = f32[] constant(1)
  %constant_577 = f32[] constant(0.99)
  %param_0.750 = s32[] parameter(0)
  %constant_456_clone_1 = s32[] constant(2147483647)
  %compare.132.clone.1 = pred[] compare(s32[] %param_0.750, s32[] %constant_456_clone_1), direction=LT, metadata={op_name="jit(training_step)/jit(main)/lt" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/numerics.py" source_line=68}
  %constant_579_clone_1 = s32[] constant(1)
  %add.408.clone.1 = s32[] add(s32[] %param_0.750, s32[] %constant_579_clone_1), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/numerics.py" source_line=68}
  %select.125.clone.1 = s32[] select(pred[] %compare.132.clone.1, s32[] %add.408.clone.1, s32[] %constant_456_clone_1), metadata={op_name="jit(training_step)/jit(main)/jit(_where)/jit(jit__where)/select_n" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/numerics.py" source_line=68}
  %convert.22 = f32[] convert(s32[] %select.125.clone.1), metadata={op_name="jit(training_step)/jit(main)/convert_element_type[new_dtype=float32 weak_type=True]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=86}
  %power.3 = f32[] power(f32[] %constant_577, f32[] %convert.22), metadata={op_name="jit(training_step)/jit(main)/pow" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=86}
  %subtract.64 = f32[] subtract(f32[] %constant_578, f32[] %power.3), metadata={op_name="jit(training_step)/jit(main)/sub" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=86}
  %constant_575_clone_1 = f32[] constant(0.9)
  %power.4.clone.1 = f32[] power(f32[] %constant_575_clone_1, f32[] %convert.22), metadata={op_name="jit(training_step)/jit(main)/pow" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=86}
  %subtract.65.clone.1 = f32[] subtract(f32[] %constant_578, f32[] %power.4.clone.1), metadata={op_name="jit(training_step)/jit(main)/sub" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=86}
  ROOT %tuple.71 = (f32[], f32[], s32[]) tuple(f32[] %subtract.64, f32[] %subtract.65.clone.1, s32[] %select.125.clone.1)
}

%region_30.1729 (Arg_0.1730: f32[], Arg_1.1731: f32[]) -> f32[] {
  %Arg_0.1730 = f32[] parameter(0)
  %Arg_1.1731 = f32[] parameter(1)
  ROOT %add.1732 = f32[] add(f32[] %Arg_0.1730, f32[] %Arg_1.1731), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
}

%fused_computation.61 (param_0.395: f32[135168,27,4], param_1.521: f32[135168,150], param_2.463: f32[2]) -> f32[135168,27] {
  %param_1.521 = f32[135168,150]{1,0} parameter(1)
  %slice.309 = f32[135168,27]{1,0} slice(f32[135168,150]{1,0} %param_1.521), slice={[0:135168], [0:27]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/slice[start_indices=(0, 0) limit_indices=(135168, 27) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=73}
  %slice.308 = f32[135168,108]{1,0} slice(f32[135168,150]{1,0} %param_1.521), slice={[0:135168], [30:138]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/slice[start_indices=(0, 30) limit_indices=(135168, 138) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=73}
  %bitcast.84 = f32[135168,27,4]{2,1,0} bitcast(f32[135168,108]{1,0} %slice.308), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reshape[new_sizes=(135168, 27, 4) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=37}
  %param_0.395 = f32[135168,27,4]{2,1,0} parameter(0)
  %cosine.1 = f32[135168,27,4]{2,1,0} cosine(f32[135168,27,4]{2,1,0} %param_0.395), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/cos" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=30}
  %multiply.354 = f32[135168,27,4]{2,1,0} multiply(f32[135168,27,4]{2,1,0} %bitcast.84, f32[135168,27,4]{2,1,0} %cosine.1), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=30}
  %slice.307 = f32[135168,27,2]{2,1,0} slice(f32[135168,27,4]{2,1,0} %multiply.354), slice={[0:135168], [0:27], [0:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/slice[start_indices=(0, 0, 0) limit_indices=(135168, 27, 2) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=31}
  %slice.306 = f32[135168,27,2]{2,1,0} slice(f32[135168,27,4]{2,1,0} %multiply.354), slice={[0:135168], [0:27], [2:4]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/slice[start_indices=(0, 0, 2) limit_indices=(135168, 27, 4) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=31}
  %add.411 = f32[135168,27,2]{2,1,0} add(f32[135168,27,2]{2,1,0} %slice.307, f32[135168,27,2]{2,1,0} %slice.306), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=32}
  %param_2.463 = f32[2]{0} parameter(2)
  %broadcast.1001 = f32[135168,27,2]{2,1,0} broadcast(f32[2]{0} %param_2.463), dimensions={2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %multiply.353 = f32[135168,27,2]{2,1,0} multiply(f32[135168,27,2]{2,1,0} %add.411, f32[135168,27,2]{2,1,0} %broadcast.1001), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %constant_585 = f32[] constant(0)
  %reduce.43 = f32[135168,27]{1,0} reduce(f32[135168,27,2]{2,1,0} %multiply.353, f32[] %constant_585), dimensions={2}, to_apply=%region_30.1729, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  ROOT %add.410 = f32[135168,27]{1,0} add(f32[135168,27]{1,0} %slice.309, f32[135168,27]{1,0} %reduce.43), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
}

%fused_computation.62 (param_0.143: f32[135168,128], param_1.522: f32[135168,128]) -> f32[135168,128] {
  %param_1.522 = f32[135168,128]{1,0} parameter(1)
  %constant_586 = f32[] constant(0)
  %broadcast.1003 = f32[135168,128]{1,0} broadcast(f32[] %constant_586), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gt" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  %compare.133 = pred[135168,128]{1,0} compare(f32[135168,128]{1,0} %param_1.522, f32[135168,128]{1,0} %broadcast.1003), direction=GT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gt" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  %param_0.143 = f32[135168,128]{1,0} parameter(0)
  ROOT %select.126 = f32[135168,128]{1,0} select(pred[135168,128]{1,0} %compare.133, f32[135168,128]{1,0} %param_0.143, f32[135168,128]{1,0} %broadcast.1003), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
}

%fused_computation.63 (param_0.145: f32[135168,128], param_1.523: f32[135168,128]) -> f32[135168,128] {
  %param_1.523 = f32[135168,128]{1,0} parameter(1)
  %constant_587 = f32[] constant(0)
  %broadcast.1005 = f32[135168,128]{1,0} broadcast(f32[] %constant_587), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gt" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  %compare.134 = pred[135168,128]{1,0} compare(f32[135168,128]{1,0} %param_1.523, f32[135168,128]{1,0} %broadcast.1005), direction=GT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gt" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=105}
  %param_0.145 = f32[135168,128]{1,0} parameter(0)
  ROOT %select.127 = f32[135168,128]{1,0} select(pred[135168,128]{1,0} %compare.134, f32[135168,128]{1,0} %param_0.145, f32[135168,128]{1,0} %broadcast.1005), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=105}
}

%fused_computation.64 (param_0.699: pred[4096,33,3], param_1.907: s32[4096,33,2], param_2.932: f32[4096,221], param_3.831: f32[135168,3], param_4.604: f32[4096,3], param_5.658: f32[4096], param_6.626: f32[4096,221]) -> f32[135168,3] {
  %param_0.699 = pred[4096,33,3]{1,0,2} parameter(0)
  %param_4.604 = f32[4096,3]{1,0} parameter(4)
  %constant_589 = f32[] constant(1)
  %broadcast.2282 = f32[4096,1]{1,0} broadcast(f32[] %constant_589), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_6.626 = f32[4096,221]{1,0} parameter(6)
  %slice.710 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_6.626), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(1,), start_index_map=(1,)) slice_sizes=(4096, 1) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=482}
  %subtract.394 = f32[4096,1]{1,0} subtract(f32[4096,1]{1,0} %broadcast.2282, f32[4096,1]{1,0} %slice.710), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %constant_1828 = f32[] constant(1e-08)
  %broadcast.2281 = f32[4096,1]{1,0} broadcast(f32[] %constant_1828), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.1024 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %subtract.394, f32[4096,1]{1,0} %broadcast.2281), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %bitcast.558 = f32[4096]{0} bitcast(f32[4096,1]{1,0} %add.1024), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_5.658 = f32[4096]{0} parameter(5)
  %broadcast.2280 = f32[4096]{0} broadcast(f32[] %constant_1828), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.1023 = f32[4096]{0} add(f32[4096]{0} %param_5.658, f32[4096]{0} %broadcast.2280), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  %divide.62 = f32[4096]{0} divide(f32[4096]{0} %bitcast.558, f32[4096]{0} %add.1023), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %constant_1743 = f32[] constant(0.000162760422)
  %broadcast.2211 = f32[4096]{0} broadcast(f32[] %constant_1743), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.585 = f32[4096]{0} multiply(f32[4096]{0} %divide.62, f32[4096]{0} %broadcast.2211), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %broadcast.2210 = f32[4096,3]{1,0} broadcast(f32[4096]{0} %multiply.585), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.584 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %param_4.604, f32[4096,3]{1,0} %broadcast.2210), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %broadcast.2209 = f32[4096,221,3]{1,0,2} broadcast(f32[4096,3]{1,0} %multiply.584), dimensions={0,2}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/broadcast_in_dim[shape=(4096, 221, 3) broadcast_dimensions=(0, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %param_2.932 = f32[4096,221]{1,0} parameter(2)
  %broadcast.1007 = f32[4096,221,3]{1,0,2} broadcast(f32[4096,221]{1,0} %param_2.932), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.357 = f32[4096,221,3]{1,0,2} multiply(f32[4096,221,3]{1,0,2} %broadcast.2209, f32[4096,221,3]{1,0,2} %broadcast.1007), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %param_1.907 = s32[4096,33,2]{1,0,2} parameter(1)
  %gather.48 = f32[4096,33,3]{1,0,2} gather(f32[4096,221,3]{1,0,2} %multiply.357, s32[4096,33,2]{1,0,2} %param_1.907), offset_dims={2}, collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=2, slice_sizes={1,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 1), start_index_map=(0, 1)) slice_sizes=(1, 1, 3) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %constant_588 = f32[] constant(0)
  %broadcast.1009 = f32[4096,33,3]{1,0,2} broadcast(f32[] %constant_588), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 33, 3) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %select.128 = f32[4096,33,3]{1,0,2} select(pred[4096,33,3]{1,0,2} %param_0.699, f32[4096,33,3]{1,0,2} %gather.48, f32[4096,33,3]{1,0,2} %broadcast.1009), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %bitcast.85 = f32[135168,3]{0,1} bitcast(f32[4096,33,3]{1,0,2} %select.128), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reshape[new_sizes=(135168, 3) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=440}
  %broadcast.1011 = f32[135168,3]{0,1} broadcast(f32[] %constant_589), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %param_3.831 = f32[135168,3]{0,1} parameter(3)
  %negate.43 = f32[135168,3]{0,1} negate(f32[135168,3]{0,1} %param_3.831), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/neg" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %exponential.24 = f32[135168,3]{0,1} exponential(f32[135168,3]{0,1} %negate.43), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/exp" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %add.933 = f32[135168,3]{0,1} add(f32[135168,3]{0,1} %exponential.24, f32[135168,3]{0,1} %broadcast.1011), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %divide.57 = f32[135168,3]{0,1} divide(f32[135168,3]{0,1} %broadcast.1011, f32[135168,3]{0,1} %add.933), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/div" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %subtract.66 = f32[135168,3]{0,1} subtract(f32[135168,3]{0,1} %broadcast.1011, f32[135168,3]{0,1} %divide.57), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/sub" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %multiply.356 = f32[135168,3]{0,1} multiply(f32[135168,3]{0,1} %bitcast.85, f32[135168,3]{0,1} %subtract.66), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(sigmoid)))/jit(jit_transpose(jvp(sigmoid)))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  ROOT %multiply.355 = f32[135168,3]{0,1} multiply(f32[135168,3]{0,1} %multiply.356, f32[135168,3]{0,1} %divide.57), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(sigmoid)))/jit(jit_transpose(jvp(sigmoid)))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
}

%fused_computation.66 (param_0.151: f32[4096,3], param_1.1002: f32[3,4096], param_2.1050: f32[4096,221], param_3.927: f32[4096]) -> f32[4096,3] {
  %param_1.1002 = f32[3,4096]{0,1} parameter(1)
  %bitcast.138 = f32[4096,3]{1,0} bitcast(f32[3,4096]{0,1} %param_1.1002), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %constant_1824 = f32[] constant(1)
  %broadcast.2276 = f32[4096,1]{1,0} broadcast(f32[] %constant_1824), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_2.1050 = f32[4096,221]{1,0} parameter(2)
  %slice.708 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %param_2.1050), slice={[0:4096], [220:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(1,), start_index_map=(1,)) slice_sizes=(4096, 1) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=482}
  %subtract.392 = f32[4096,1]{1,0} subtract(f32[4096,1]{1,0} %broadcast.2276, f32[4096,1]{1,0} %slice.708), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %constant_1823 = f32[] constant(1e-08)
  %broadcast.2275 = f32[4096,1]{1,0} broadcast(f32[] %constant_1823), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.1020 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %subtract.392, f32[4096,1]{1,0} %broadcast.2275), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %bitcast.556 = f32[4096]{0} bitcast(f32[4096,1]{1,0} %add.1020), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %param_3.927 = f32[4096]{0} parameter(3)
  %broadcast.2274 = f32[4096]{0} broadcast(f32[] %constant_1823), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=481}
  %add.1019 = f32[4096]{0} add(f32[4096]{0} %param_3.927, f32[4096]{0} %broadcast.2274), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  %divide.60 = f32[4096]{0} divide(f32[4096]{0} %bitcast.556, f32[4096]{0} %add.1019), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=479}
  %broadcast.898 = f32[4096,3]{1,0} broadcast(f32[4096]{0} %divide.60), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.360 = f32[4096,3]{1,0} multiply(f32[4096,3]{1,0} %bitcast.138, f32[4096,3]{1,0} %broadcast.898), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %bitcast.86 = f32[4096]{0} bitcast(f32[4096,1]{1,0} %slice.708), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=196}
  %broadcast.897 = f32[4096,3]{1,0} broadcast(f32[4096]{0} %bitcast.86), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=196}
  %add.412 = f32[4096,3]{1,0} add(f32[4096,3]{1,0} %multiply.360, f32[4096,3]{1,0} %broadcast.897), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=196}
  %param_0.151 = f32[4096,3]{1,0} parameter(0)
  ROOT %subtract.67 = f32[4096,3]{1,0} subtract(f32[4096,3]{1,0} %add.412, f32[4096,3]{1,0} %param_0.151), metadata={op_name="jit(training_step)/jit(main)/sub" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
}

%region_8.953 (Arg_0.954: f32[], Arg_1.955: f32[]) -> f32[] {
  %Arg_0.954 = f32[] parameter(0)
  %Arg_1.955 = f32[] parameter(1)
  ROOT %add.956 = f32[] add(f32[] %Arg_0.954, f32[] %Arg_1.955), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
}

%input_fused_computation_reduce.4 (param_0.396: f32[4096,221], param_1.526: s32[4096,33,2]) -> f32[4096] {
  %param_0.396 = f32[4096,221]{1,0} parameter(0)
  %param_1.526 = s32[4096,33,2]{1,0,2} parameter(1)
  %gather.49 = f32[4096,33]{1,0} gather(f32[4096,221]{1,0} %param_0.396, s32[4096,33,2]{1,0,2} %param_1.526), offset_dims={}, collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=2, slice_sizes={1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0, 1), start_index_map=(0, 1)) slice_sizes=(1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=472}
  %constant_592 = f32[] constant(0)
  ROOT %reduce.44 = f32[4096]{0} reduce(f32[4096,33]{1,0} %gather.49, f32[] %constant_592), dimensions={1}, to_apply=%region_8.953, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
}

%region_9.957 (Arg_0.958: f32[], Arg_1.959: f32[]) -> f32[] {
  %Arg_0.958 = f32[] parameter(0)
  %Arg_1.959 = f32[] parameter(1)
  ROOT %add.960 = f32[] add(f32[] %Arg_0.958, f32[] %Arg_1.959), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
}

%input_fused_computation_reduce.5 (param_0.399: f32[4096,221,3], param_1.529: f32[4096,221]) -> f32[3,4096] {
  %param_0.399 = f32[4096,221,3]{1,0,2} parameter(0)
  %param_1.529 = f32[4096,221]{1,0} parameter(1)
  %broadcast.1013 = f32[4096,221,3]{1,0,2} broadcast(f32[4096,221]{1,0} %param_1.529), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %multiply.361 = f32[4096,221,3]{1,0,2} multiply(f32[4096,221,3]{1,0,2} %param_0.399, f32[4096,221,3]{1,0,2} %broadcast.1013), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %bitcast.88 = f32[3,4096,221]{2,1,0} bitcast(f32[4096,221,3]{1,0,2} %multiply.361), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %constant_595 = f32[] constant(0)
  ROOT %reduce.45 = f32[3,4096]{0,1} reduce(f32[3,4096,221]{2,1,0} %bitcast.88, f32[] %constant_595), dimensions={2}, to_apply=%region_9.957
}

%region_7.949 (Arg_0.950: f32[], Arg_1.951: f32[]) -> f32[] {
  %Arg_0.950 = f32[] parameter(0)
  %Arg_1.951 = f32[] parameter(1)
  ROOT %add.952 = f32[] add(f32[] %Arg_0.950, f32[] %Arg_1.951), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%input_fused_computation_scatter.13 (param_0.161: s32[4096,33,2], param_1.240: pred[4096,33,3], param_2.911: f32[135168,3]) -> f32[4096,221,3] {
  %constant_596 = f32[] constant(0)
  %broadcast.902 = f32[4096,221,3]{1,0,2} broadcast(f32[] %constant_596), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 221, 3) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %param_0.161 = s32[4096,33,2]{1,0,2} parameter(0)
  %param_1.240 = pred[4096,33,3]{1,0,2} parameter(1)
  %constant_1705 = f32[] constant(1)
  %broadcast.2167 = f32[135168,3]{0,1} broadcast(f32[] %constant_1705), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %param_2.911 = f32[135168,3]{0,1} parameter(2)
  %negate.41 = f32[135168,3]{0,1} negate(f32[135168,3]{0,1} %param_2.911), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/neg" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %exponential.22 = f32[135168,3]{0,1} exponential(f32[135168,3]{0,1} %negate.41), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/exp" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %add.931 = f32[135168,3]{0,1} add(f32[135168,3]{0,1} %exponential.22, f32[135168,3]{0,1} %broadcast.2167), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %divide.55 = f32[135168,3]{0,1} divide(f32[135168,3]{0,1} %broadcast.2167, f32[135168,3]{0,1} %add.931), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(sigmoid))/jit(jit_jvp(sigmoid))/div" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %bitcast.89 = f32[4096,33,3]{1,0,2} bitcast(f32[135168,3]{0,1} %divide.55), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 3) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=440}
  %broadcast.1015 = f32[4096,33,3]{1,0,2} broadcast(f32[] %constant_596), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 33, 3) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %select.129 = f32[4096,33,3]{1,0,2} select(pred[4096,33,3]{1,0,2} %param_1.240, f32[4096,33,3]{1,0,2} %bitcast.89, f32[4096,33,3]{1,0,2} %broadcast.1015), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  ROOT %scatter.55 = f32[4096,221,3]{1,0,2} scatter(f32[4096,221,3]{1,0,2} %broadcast.902, s32[4096,33,2]{1,0,2} %param_0.161, f32[4096,33,3]{1,0,2} %select.129), update_window_dims={2}, inserted_window_dims={0,1}, scatter_dims_to_operand_dims={0,1}, index_vector_dim=2, to_apply=%region_7.949, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 1), scatter_dims_to_operand_dims=(0, 1)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%fused_computation.70 (param_0.700: u32[4096,221,3], param_1.912: s32[4096,33,2]) -> pred[4096,33,3] {
  %iota.208 = u32[4096,33,1]{2,1,0} iota(), iota_dimension=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/iota[dtype=uint32 shape=(135168,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %iota.207 = u32[4096,33,1]{2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/iota[dtype=uint32 shape=(135168,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %constant_1718 = u32[] constant(33), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %broadcast.2185 = u32[4096,33,1]{2,1,0} broadcast(u32[] %constant_1718), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %multiply.577 = u32[4096,33,1]{2,1,0} multiply(u32[4096,33,1]{2,1,0} %iota.207, u32[4096,33,1]{2,1,0} %broadcast.2185), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %add.941 = u32[4096,33,1]{2,1,0} add(u32[4096,33,1]{2,1,0} %iota.208, u32[4096,33,1]{2,1,0} %multiply.577), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %constant_1717 = u32[] constant(1)
  %broadcast.2184 = u32[4096,33,1]{2,1,0} broadcast(u32[] %constant_1717), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 33, 3) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %add.940 = u32[4096,33,1]{2,1,0} add(u32[4096,33,1]{2,1,0} %add.941, u32[4096,33,1]{2,1,0} %broadcast.2184), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %bitcast.546 = u32[4096,33]{1,0} bitcast(u32[4096,33,1]{2,1,0} %add.940), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %broadcast.2183 = u32[4096,33,3]{1,0,2} broadcast(u32[4096,33]{1,0} %bitcast.546), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %param_0.700 = u32[4096,221,3]{1,0,2} parameter(0)
  %param_1.912 = s32[4096,33,2]{1,0,2} parameter(1)
  %gather.50 = u32[4096,33,3]{1,0,2} gather(u32[4096,221,3]{1,0,2} %param_0.700, s32[4096,33,2]{1,0,2} %param_1.912), offset_dims={2}, collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=2, slice_sizes={1,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 1), start_index_map=(0, 1)) slice_sizes=(1, 1, 3) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  ROOT %compare.135 = pred[4096,33,3]{1,0,2} compare(u32[4096,33,3]{1,0,2} %broadcast.2183, u32[4096,33,3]{1,0,2} %gather.50), direction=EQ, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%region_6.946 (Arg_0.947: u32[], Arg_1.948: u32[]) -> u32[] {
  %Arg_0.947 = u32[] parameter(0)
  ROOT %Arg_1.948 = u32[] parameter(1)
}

%input_fused_computation_scatter.14 (param_0.165: s32[4096,33,2]) -> u32[4096,221,3] {
  %constant_457 = u32[] constant(0)
  %broadcast.903 = u32[4096,221,3]{1,0,2} broadcast(u32[] %constant_457), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 221, 3) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %param_0.165 = s32[4096,33,2]{1,0,2} parameter(0)
  %iota.204 = u32[4096,33,1]{2,1,0} iota(), iota_dimension=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/iota[dtype=uint32 shape=(135168,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %iota.202 = u32[4096,33,1]{2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/iota[dtype=uint32 shape=(135168,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %constant_1713 = u32[] constant(33), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %broadcast.2179 = u32[4096,33,1]{2,1,0} broadcast(u32[] %constant_1713), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %multiply.575 = u32[4096,33,1]{2,1,0} multiply(u32[4096,33,1]{2,1,0} %iota.202, u32[4096,33,1]{2,1,0} %broadcast.2179), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %add.937 = u32[4096,33,1]{2,1,0} add(u32[4096,33,1]{2,1,0} %iota.204, u32[4096,33,1]{2,1,0} %multiply.575), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(4096, 33, 1) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %constant_1712 = u32[] constant(1)
  %broadcast.2178 = u32[4096,33,1]{2,1,0} broadcast(u32[] %constant_1712), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 33, 3) broadcast_dimensions=()]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %add.936 = u32[4096,33,1]{2,1,0} add(u32[4096,33,1]{2,1,0} %add.937, u32[4096,33,1]{2,1,0} %broadcast.2178), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %bitcast.544 = u32[4096,33]{1,0} bitcast(u32[4096,33,1]{2,1,0} %add.936), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %broadcast.2177 = u32[4096,33,3]{1,0,2} broadcast(u32[4096,33]{1,0} %bitcast.544), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  ROOT %scatter.56 = u32[4096,221,3]{1,0,2} scatter(u32[4096,221,3]{1,0,2} %broadcast.903, s32[4096,33,2]{1,0,2} %param_0.165, u32[4096,33,3]{1,0,2} %broadcast.2177), update_window_dims={2}, inserted_window_dims={0,1}, scatter_dims_to_operand_dims={0,1}, index_vector_dim=2, to_apply=%region_6.946, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/scatter[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 1), scatter_dims_to_operand_dims=(0, 1)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
}

%fused_computation.72 (param_0.170: f32[135168,27], param_1.535: f32[2], param_2.472: f32[135168,27,4], param_3.366: f32[4096,3]) -> f32[135168,150] {
  %param_0.170 = f32[135168,27]{1,0} parameter(0)
  %param_3.366 = f32[4096,3]{1,0} parameter(3)
  %broadcast.910 = f32[1,4096,33,1,1,3]{5,4,3,2,1,0} broadcast(f32[4096,3]{1,0} %param_3.366), dimensions={1,5}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(1, 4096, 33, 1, 1, 3) broadcast_dimensions=(0, 1, 3, 4, 5)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=435}
  %bitcast.93 = f32[135168,3]{1,0} bitcast(f32[1,4096,33,1,1,3]{5,4,3,2,1,0} %broadcast.910), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(135168, 3) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=435}
  %param_2.472 = f32[135168,27,4]{2,1,0} parameter(2)
  %sine.3 = f32[135168,27,4]{2,1,0} sine(f32[135168,27,4]{2,1,0} %param_2.472), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sin" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=30}
  %bitcast.92 = f32[135168,108]{1,0} bitcast(f32[135168,27,4]{2,1,0} %sine.3), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(135168, 108) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=37}
  %broadcast.909 = f32[135168,3,2]{2,1,0} broadcast(f32[135168,3]{1,0} %bitcast.93), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %param_1.535 = f32[2]{0} parameter(1)
  %broadcast.908 = f32[135168,3,2]{2,1,0} broadcast(f32[2]{0} %param_1.535), dimensions={2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %multiply.363 = f32[135168,3,2]{2,1,0} multiply(f32[135168,3,2]{2,1,0} %broadcast.909, f32[135168,3,2]{2,1,0} %broadcast.908), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %constant_600 = f32[] constant(1.57079637)
  %broadcast.907 = f32[135168,3,2]{2,1,0} broadcast(f32[] %constant_600), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=32}
  %add.420 = f32[135168,3,2]{2,1,0} add(f32[135168,3,2]{2,1,0} %multiply.363, f32[135168,3,2]{2,1,0} %broadcast.907), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=32}
  %concatenate.109 = f32[135168,3,4]{2,1,0} concatenate(f32[135168,3,2]{2,1,0} %multiply.363, f32[135168,3,2]{2,1,0} %add.420), dimensions={2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=2]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=31}
  %sine.2 = f32[135168,3,4]{2,1,0} sine(f32[135168,3,4]{2,1,0} %concatenate.109), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sin" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=30}
  %bitcast.91 = f32[135168,12]{1,0} bitcast(f32[135168,3,4]{2,1,0} %sine.2), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(135168, 12) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=37}
  ROOT %concatenate.108 = f32[135168,150]{1,0} concatenate(f32[135168,27]{1,0} %param_0.170, f32[135168,3]{1,0} %bitcast.93, f32[135168,108]{1,0} %bitcast.92, f32[135168,12]{1,0} %bitcast.91), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=73}
}

%fused_computation.73 (param_0.404: f32[135168,27], param_1.537: f32[2]) -> f32[135168,27,4] {
  %param_0.404 = f32[135168,27]{1,0} parameter(0)
  %broadcast.912 = f32[135168,27,2]{2,1,0} broadcast(f32[135168,27]{1,0} %param_0.404), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %param_1.537 = f32[2]{0} parameter(1)
  %broadcast.1022 = f32[135168,27,2]{2,1,0} broadcast(f32[2]{0} %param_1.537), dimensions={2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %multiply.364 = f32[135168,27,2]{2,1,0} multiply(f32[135168,27,2]{2,1,0} %broadcast.912, f32[135168,27,2]{2,1,0} %broadcast.1022), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %constant_601 = f32[] constant(1.57079637)
  %broadcast.911 = f32[135168,27,2]{2,1,0} broadcast(f32[] %constant_601), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=32}
  %add.421 = f32[135168,27,2]{2,1,0} add(f32[135168,27,2]{2,1,0} %multiply.364, f32[135168,27,2]{2,1,0} %broadcast.911), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=32}
  ROOT %concatenate.110 = f32[135168,27,4]{2,1,0} concatenate(f32[135168,27,2]{2,1,0} %multiply.364, f32[135168,27,2]{2,1,0} %add.421), dimensions={2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=2]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=31}
}

%fused_computation.75 (param_0.696: f32[3,48,128], param_1.906: f32[3,4096,3,33], param_2.909: f32[3,48,128,128]) -> f32[135168,144] {
  %constant_1357 = f32[] constant(1)
  %broadcast.1790 = f32[3,4096,33]{2,1,0} broadcast(f32[] %constant_1357), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.906 = f32[3,4096,3,33]{3,1,0,2} parameter(1)
  %slice.450 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.906), slice={[0:3], [0:4096], [0:1], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1, 2, 3), collapsed_slice_dims=(), start_index_map=(2,)) slice_sizes=(3, 4096, 1, 33) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=165}
  %bitcast.302 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.450), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.251 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.302), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.272 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.302, f32[3,4096,33]{2,1,0} %floor.251), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.271 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1790, f32[3,4096,33]{2,1,0} %subtract.272), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1789 = f32[3,4096,48,33]{2,1,0,3} broadcast(f32[3,4096,33]{2,1,0} %subtract.271), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.696 = f32[3,48,128]{2,1,0} parameter(0)
  %iota.167 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1353 = s32[] constant(0)
  %broadcast.1788 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1353), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.146 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.251), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1351 = s32[] constant(127)
  %broadcast.1787 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1351), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.135 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1788, s32[3,4096,33]{2,1,0} %convert.146, s32[3,4096,33]{2,1,0} %broadcast.1787), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.261 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.135, s32[3,4096,33]{2,1,0} %broadcast.1788), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1350 = s32[] constant(128)
  %broadcast.1786 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1350), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.781 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.135, s32[3,4096,33]{2,1,0} %broadcast.1786), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.309 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.261, s32[3,4096,33]{2,1,0} %add.781, s32[3,4096,33]{2,1,0} %clamp.135), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.300 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.309), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.219 = s32[3,4096,33,2]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.167, s32[3,4096,33,1]{2,1,0,3} %bitcast.300), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.111 = f32[3,4096,48,33]{2,1,0,3} gather(f32[3,48,128]{2,1,0} %param_0.696, s32[3,4096,33,2]{2,1,0,3} %concatenate.219), offset_dims={2}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,48,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 48, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.524 = f32[3,4096,48,33]{2,1,0,3} multiply(f32[3,4096,48,33]{2,1,0,3} %broadcast.1789, f32[3,4096,48,33]{2,1,0,3} %gather.111), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.39 = f32[3,4096,48,33]{3,2,1,0} copy(f32[3,4096,48,33]{2,1,0,3} %multiply.524), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1785 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %subtract.272), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.166 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_1352 = s32[] constant(1)
  %broadcast.1781 = s32[3,4096,33]{2,1,0} broadcast(s32[] %constant_1352), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.780 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.146, s32[3,4096,33]{2,1,0} %broadcast.1781), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.134 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1788, s32[3,4096,33]{2,1,0} %add.780, s32[3,4096,33]{2,1,0} %broadcast.1787), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.260 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.134, s32[3,4096,33]{2,1,0} %broadcast.1788), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.779 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.134, s32[3,4096,33]{2,1,0} %broadcast.1786), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.308 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.260, s32[3,4096,33]{2,1,0} %add.779, s32[3,4096,33]{2,1,0} %clamp.134), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.297 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.308), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.218 = s32[3,4096,33,2]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.166, s32[3,4096,33,1]{3,2,1,0} %bitcast.297), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.110 = f32[3,4096,48,33]{3,2,1,0} gather(f32[3,48,128]{2,1,0} %param_0.696, s32[3,4096,33,2]{3,2,1,0} %concatenate.218), offset_dims={2}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,48,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 48, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.523 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.1785, f32[3,4096,48,33]{3,2,1,0} %gather.110), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.778 = f32[3,4096,48,33]{3,2,1,0} add(f32[3,4096,48,33]{3,2,1,0} %copy.39, f32[3,4096,48,33]{3,2,1,0} %multiply.523), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.650 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.906), slice={[0:3], [0:4096], [1:2], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 4096, 1, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.542 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.650), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.398 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.542), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.374 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.542, f32[3,4096,33]{2,1,0} %floor.398), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.373 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1790, f32[3,4096,33]{2,1,0} %subtract.374), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.649 = f32[3,4096,1,33]{3,1,0,2} slice(f32[3,4096,3,33]{3,1,0,2} %param_1.906), slice={[0:3], [0:4096], [2:3], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 1, 0) limit_indices=(3, 4096, 2, 33) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.541 = f32[3,4096,33]{2,1,0} bitcast(f32[3,4096,1,33]{3,1,0,2} %slice.649), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %floor.397 = f32[3,4096,33]{2,1,0} floor(f32[3,4096,33]{2,1,0} %bitcast.541), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.370 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %bitcast.541, f32[3,4096,33]{2,1,0} %floor.397), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.369 = f32[3,4096,33]{2,1,0} subtract(f32[3,4096,33]{2,1,0} %broadcast.1790, f32[3,4096,33]{2,1,0} %subtract.370), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.573 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.373, f32[3,4096,33]{2,1,0} %subtract.369), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2161 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.573), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.909 = f32[3,48,128,128]{3,2,1,0} parameter(2)
  %iota.199 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.226 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.398), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.215 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1788, s32[3,4096,33]{2,1,0} %convert.226, s32[3,4096,33]{2,1,0} %broadcast.1787), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.343 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.215, s32[3,4096,33]{2,1,0} %broadcast.1788), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.929 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.215, s32[3,4096,33]{2,1,0} %broadcast.1786), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.391 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.343, s32[3,4096,33]{2,1,0} %add.929, s32[3,4096,33]{2,1,0} %clamp.215), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.539 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.391), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.225 = s32[3,4096,33]{2,1,0} convert(f32[3,4096,33]{2,1,0} %floor.397), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.214 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1788, s32[3,4096,33]{2,1,0} %convert.225, s32[3,4096,33]{2,1,0} %broadcast.1787), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.342 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.214, s32[3,4096,33]{2,1,0} %broadcast.1788), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.928 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.214, s32[3,4096,33]{2,1,0} %broadcast.1786), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.390 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.342, s32[3,4096,33]{2,1,0} %add.928, s32[3,4096,33]{2,1,0} %clamp.214), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.537 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.390), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.251 = s32[3,4096,33,3]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.199, s32[3,4096,33,1]{2,1,0,3} %bitcast.539, s32[3,4096,33,1]{2,1,0,3} %bitcast.537), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.127 = f32[3,4096,48,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.909, s32[3,4096,33,3]{2,1,0,3} %concatenate.251), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.572 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.2161, f32[3,4096,48,33]{3,2,1,0} %gather.127), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.571 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.373, f32[3,4096,33]{2,1,0} %subtract.370), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2148 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.571), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.198 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.533 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.391), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.926 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.225, s32[3,4096,33]{2,1,0} %broadcast.1781), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.212 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1788, s32[3,4096,33]{2,1,0} %add.926, s32[3,4096,33]{2,1,0} %broadcast.1787), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.340 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.212, s32[3,4096,33]{2,1,0} %broadcast.1788), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.925 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.212, s32[3,4096,33]{2,1,0} %broadcast.1786), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.388 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.340, s32[3,4096,33]{2,1,0} %add.925, s32[3,4096,33]{2,1,0} %clamp.212), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.531 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.388), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.250 = s32[3,4096,33,3]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.198, s32[3,4096,33,1]{3,2,1,0} %bitcast.533, s32[3,4096,33,1]{3,2,1,0} %bitcast.531), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.126 = f32[3,4096,48,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.909, s32[3,4096,33,3]{3,2,1,0} %concatenate.250), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.570 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.2148, f32[3,4096,48,33]{3,2,1,0} %gather.126), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.924 = f32[3,4096,48,33]{3,2,1,0} add(f32[3,4096,48,33]{3,2,1,0} %multiply.572, f32[3,4096,48,33]{3,2,1,0} %multiply.570), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.569 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.374, f32[3,4096,33]{2,1,0} %subtract.369), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2130 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.569), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.197 = s32[3,4096,33,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.923 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %convert.226, s32[3,4096,33]{2,1,0} %broadcast.1781), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.211 = s32[3,4096,33]{2,1,0} clamp(s32[3,4096,33]{2,1,0} %broadcast.1788, s32[3,4096,33]{2,1,0} %add.923, s32[3,4096,33]{2,1,0} %broadcast.1787), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.339 = pred[3,4096,33]{2,1,0} compare(s32[3,4096,33]{2,1,0} %clamp.211, s32[3,4096,33]{2,1,0} %broadcast.1788), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.922 = s32[3,4096,33]{2,1,0} add(s32[3,4096,33]{2,1,0} %clamp.211, s32[3,4096,33]{2,1,0} %broadcast.1786), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.387 = s32[3,4096,33]{2,1,0} select(pred[3,4096,33]{2,1,0} %compare.339, s32[3,4096,33]{2,1,0} %add.922, s32[3,4096,33]{2,1,0} %clamp.211), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.527 = s32[3,4096,33,1]{2,1,0,3} bitcast(s32[3,4096,33]{2,1,0} %select.387), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.249 = s32[3,4096,33,3]{2,1,0,3} concatenate(s32[3,4096,33,1]{2,1,0,3} %iota.197, s32[3,4096,33,1]{2,1,0,3} %bitcast.527, s32[3,4096,33,1]{2,1,0,3} %bitcast.537), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.125 = f32[3,4096,48,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.909, s32[3,4096,33,3]{2,1,0,3} %concatenate.249), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.568 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.2130, f32[3,4096,48,33]{3,2,1,0} %gather.125), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.920 = f32[3,4096,48,33]{3,2,1,0} add(f32[3,4096,48,33]{3,2,1,0} %add.924, f32[3,4096,48,33]{3,2,1,0} %multiply.568), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.567 = f32[3,4096,33]{2,1,0} multiply(f32[3,4096,33]{2,1,0} %subtract.374, f32[3,4096,33]{2,1,0} %subtract.370), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.2121 = f32[3,4096,48,33]{3,2,1,0} broadcast(f32[3,4096,33]{2,1,0} %multiply.567), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.196 = s32[3,4096,33,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 33, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.521 = s32[3,4096,33,1]{3,2,1,0} bitcast(s32[3,4096,33]{2,1,0} %select.387), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 33, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.248 = s32[3,4096,33,3]{3,2,1,0} concatenate(s32[3,4096,33,1]{3,2,1,0} %iota.196, s32[3,4096,33,1]{3,2,1,0} %bitcast.521, s32[3,4096,33,1]{3,2,1,0} %bitcast.531), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.124 = f32[3,4096,48,33]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.909, s32[3,4096,33,3]{3,2,1,0} %concatenate.248), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,48,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 48, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.566 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %broadcast.2121, f32[3,4096,48,33]{3,2,1,0} %gather.124), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.915 = f32[3,4096,48,33]{3,2,1,0} add(f32[3,4096,48,33]{3,2,1,0} %add.920, f32[3,4096,48,33]{3,2,1,0} %multiply.566), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.365 = f32[3,4096,48,33]{3,2,1,0} multiply(f32[3,4096,48,33]{3,2,1,0} %add.778, f32[3,4096,48,33]{3,2,1,0} %add.915), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %transpose.30 = f32[3,48,4096,33]{3,1,2,0} transpose(f32[3,4096,48,33]{3,2,1,0} %multiply.365), dimensions={0,2,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/transpose[permutation=(0, 2, 1, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=90}
  %copy.29 = f32[3,48,4096,33]{3,2,1,0} copy(f32[3,48,4096,33]{3,1,2,0} %transpose.30), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/transpose[permutation=(0, 2, 1, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=90}
  %bitcast.95 = f32[144,4096,33]{2,1,0} bitcast(f32[3,48,4096,33]{3,2,1,0} %copy.29), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/reshape[new_sizes=(144, 4096, 33) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=95}
  %transpose.29 = f32[4096,33,144]{1,0,2} transpose(f32[144,4096,33]{2,1,0} %bitcast.95), dimensions={1,2,0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/transpose[permutation=(1, 2, 0)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
  ROOT %bitcast.94 = f32[135168,144]{0,1} bitcast(f32[4096,33,144]{1,0,2} %transpose.29), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(135168, 144) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
}

%fused_computation.96 (param_0.431: s32[3,1], param_1.580: s32[3,1], param_2.523: f32[3,4096,221], param_3.408: s32[4096,33,2]) -> f32[3,4096,3,33] {
  %param_2.523 = f32[3,4096,221]{2,1,0} parameter(2)
  %param_3.408 = s32[4096,33,2]{1,0,2} parameter(3)
  %gather.59 = f32[3,4096,33]{2,1,0} gather(f32[3,4096,221]{2,1,0} %param_2.523, s32[4096,33,2]{1,0,2} %param_3.408), offset_dims={0}, collapsed_slice_dims={1,2}, start_index_map={1,2}, index_vector_dim=2, slice_sizes={3,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(1, 2), start_index_map=(1, 2)) slice_sizes=(3, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %bitcast.109 = f32[1,3,4096,33]{3,2,0,1} bitcast(f32[3,4096,33]{2,1,0} %gather.59), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 33) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %param_1.580 = s32[3,1]{1,0} parameter(1)
  %gather.58 = f32[3,4096,33]{2,1,0} gather(f32[3,4096,33]{2,1,0} %gather.59, s32[3,1]{1,0} %param_1.580), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,4096,33}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 4096, 33) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=58}
  %bitcast.108 = f32[1,3,4096,33]{3,2,0,1} bitcast(f32[3,4096,33]{2,1,0} %gather.58), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 33) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %param_0.431 = s32[3,1]{1,0} parameter(0)
  %gather.57 = f32[3,4096,33]{2,1,0} gather(f32[3,4096,33]{2,1,0} %gather.59, s32[3,1]{1,0} %param_0.431), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,4096,33}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 4096, 33) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=59}
  %bitcast.107 = f32[1,3,4096,33]{3,2,0,1} bitcast(f32[3,4096,33]{2,1,0} %gather.57), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 33) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %concatenate.117 = f32[3,3,4096,33]{3,2,0,1} concatenate(f32[1,3,4096,33]{3,2,0,1} %bitcast.109, f32[1,3,4096,33]{3,2,0,1} %bitcast.108, f32[1,3,4096,33]{3,2,0,1} %bitcast.107), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %constant_629 = f32[] constant(1)
  %broadcast.920 = f32[3,3,4096,33]{3,2,0,1} broadcast(f32[] %constant_629), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=159}
  %add.437 = f32[3,3,4096,33]{3,2,0,1} add(f32[3,3,4096,33]{3,2,0,1} %concatenate.117, f32[3,3,4096,33]{3,2,0,1} %broadcast.920), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=159}
  %constant_628 = f32[] constant(63.5)
  %broadcast.919 = f32[3,3,4096,33]{3,2,0,1} broadcast(f32[] %constant_628), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=162}
  %multiply.376 = f32[3,3,4096,33]{3,2,0,1} multiply(f32[3,3,4096,33]{3,2,0,1} %add.437, f32[3,3,4096,33]{3,2,0,1} %broadcast.919), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=162}
  ROOT %transpose.31 = f32[3,4096,3,33]{3,1,0,2} transpose(f32[3,3,4096,33]{3,2,0,1} %multiply.376), dimensions={0,2,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/transpose[permutation=(0, 2, 1, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=165}
}

%fused_computation.97 (param_0.432: s32[4096,221]) -> s32[4096,33,2] {
  %iota.92 = s32[4096,33,1]{1,0,2} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 33) broadcast_dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %param_0.432 = s32[4096,221]{1,0} parameter(0)
  %slice.313 = s32[4096,33]{1,0} slice(s32[4096,221]{1,0} %param_0.432), slice={[0:4096], [0:33]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(4096, 33) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_630 = s32[] constant(0)
  %broadcast.922 = s32[4096,33]{1,0} broadcast(s32[] %constant_630), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %compare.142 = pred[4096,33]{1,0} compare(s32[4096,33]{1,0} %slice.313, s32[4096,33]{1,0} %broadcast.922), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %constant_462 = s32[] constant(221)
  %broadcast.921 = s32[4096,33]{1,0} broadcast(s32[] %constant_462), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %add.438 = s32[4096,33]{1,0} add(s32[4096,33]{1,0} %slice.313, s32[4096,33]{1,0} %broadcast.921), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %select.136 = s32[4096,33]{1,0} select(pred[4096,33]{1,0} %compare.142, s32[4096,33]{1,0} %add.438, s32[4096,33]{1,0} %slice.313), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %bitcast.110 = s32[4096,33,1]{1,0,2} bitcast(s32[4096,33]{1,0} %select.136), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/broadcast_in_dim[shape=(4096, 33, 1) broadcast_dimensions=(0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  ROOT %concatenate.118 = s32[4096,33,2]{1,0,2} concatenate(s32[4096,33,1]{1,0,2} %iota.92, s32[4096,33,1]{1,0,2} %bitcast.110), dimensions={2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=2]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
}

%fused_computation.98 (param_0.751: f32[221], param_1.989: pred[4096], param_2.1033: f32[4096], param_3.905: f32[4096], param_4.669: f32[4096,221], param_5.625: f32[4096,110], param_6.596: f32[4096,110]) -> (f32[4096,221], f32[4096,221], f32[4096,221]) {
  %param_0.751 = f32[221]{0} parameter(0)
  %broadcast.923 = f32[4096,221]{1,0} broadcast(f32[221]{0} %param_0.751), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_1275_clone_1 = f32[] constant(1)
  %broadcast.1674.clone.1 = f32[4096,221]{1,0} broadcast(f32[] %constant_1275_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %param_4.669 = f32[4096,221]{1,0} parameter(4)
  %negate.33.clone.1 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_4.669), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_1.989 = pred[4096]{0} parameter(1)
  %param_3.905 = f32[4096]{0} parameter(3)
  %param_2.1033 = f32[4096]{0} parameter(2)
  %constant_1232_clone_1 = f32[] constant(0.001)
  %broadcast.1673.clone.1 = f32[4096]{0} broadcast(f32[] %constant_1232_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.709.clone.1 = f32[4096]{0} add(f32[4096]{0} %param_2.1033, f32[4096]{0} %broadcast.1673.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.83.clone.1 = f32[4096]{0} maximum(f32[4096]{0} %param_3.905, f32[4096]{0} %add.709.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.279.clone.1 = f32[4096]{0} select(pred[4096]{0} %param_1.989, f32[4096]{0} %maximum.83.clone.1, f32[4096]{0} %broadcast.1673.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %constant_1231_clone_1 = f32[] constant(0)
  %broadcast.1672.clone.1 = f32[4096]{0} broadcast(f32[] %constant_1231_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.278.clone.1 = f32[4096]{0} select(pred[4096]{0} %param_1.989, f32[4096]{0} %param_2.1033, f32[4096]{0} %broadcast.1672.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.244.clone.1 = f32[4096]{0} subtract(f32[4096]{0} %select.279.clone.1, f32[4096]{0} %select.278.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1230_clone_1 = f32[] constant(0.00452488707)
  %broadcast.1671.clone.1 = f32[4096]{0} broadcast(f32[] %constant_1230_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.503.clone.1 = f32[4096]{0} multiply(f32[4096]{0} %subtract.244.clone.1, f32[4096]{0} %broadcast.1671.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.1670.clone.1 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.503.clone.1), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.502.clone.1 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.33.clone.1, f32[4096,221]{1,0} %broadcast.1670.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %exponential.18.clone.1 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %multiply.502.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %subtract.243.clone.1 = f32[4096,221]{1,0} subtract(f32[4096,221]{1,0} %broadcast.1674.clone.1, f32[4096,221]{1,0} %exponential.18.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=276}
  %slice.316.clone.1 = f32[4096,1]{1,0} slice(f32[4096,221]{1,0} %multiply.502.clone.1), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_5.625 = f32[4096,110]{1,0} parameter(5)
  %param_6.596 = f32[4096,110]{1,0} parameter(6)
  %add.441.clone.1 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %param_5.625, f32[4096,110]{1,0} %param_6.596), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.315.clone.1 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %multiply.502.clone.1), slice={[0:4096:1], [2:221:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.440.clone.1 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %add.441.clone.1, f32[4096,110]{1,0} %slice.315.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.119.clone.1 = f32[4096,111]{1,0} concatenate(f32[4096,1]{1,0} %slice.316.clone.1, f32[4096,110]{1,0} %add.440.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.115.clone.1 = f32[4096,221]{1,0} pad(f32[4096,111]{1,0} %concatenate.119.clone.1, f32[] %constant_1231_clone_1), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.114.clone.1 = f32[4096,221]{1,0} pad(f32[4096,110]{1,0} %add.441.clone.1, f32[] %constant_1231_clone_1), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.439.clone.1 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %pad.115.clone.1, f32[4096,221]{1,0} %pad.114.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %exponential.13.clone.1 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %add.439.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.410.clone.1 = f32[4096,220]{1,0} slice(f32[4096,221]{1,0} %exponential.13.clone.1), slice={[0:4096], [0:220]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(1,)) slice_sizes=(4096, 220) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=286}
  %pad.144.clone.1 = f32[4096,221]{1,0} pad(f32[4096,220]{1,0} %slice.410.clone.1, f32[] %constant_1275_clone_1), padding=0_0x1_0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=283}
  %multiply.499.clone.1 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %subtract.243.clone.1, f32[4096,221]{1,0} %pad.144.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=280}
  %log.4 = f32[4096,221]{1,0} log(f32[4096,221]{1,0} %multiply.499.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/log" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %subtract.70 = f32[4096,221]{1,0} subtract(f32[4096,221]{1,0} %broadcast.923, f32[4096,221]{1,0} %log.4), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  ROOT %tuple.76 = (f32[4096,221]{1,0}, f32[4096,221]{1,0}, f32[4096,221]{1,0}) tuple(f32[4096,221]{1,0} %subtract.70, f32[4096,221]{1,0} %multiply.499.clone.1, f32[4096,221]{1,0} %exponential.13.clone.1)
}

%fused_computation.105 (param_0.584: f32[4096,3], param_1.799: f32[4096,13]) -> f32[4096,6] {
  %param_1.799 = f32[4096,13]{1,0} parameter(1)
  %slice.388 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_1.799), slice={[0:4096:1], [0:12:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 12) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.387 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_1.799), slice={[0:4096:1], [1:13:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 13) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.726 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %slice.388, f32[4096,6]{1,0} %slice.387), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.328 = f32[4096,1]{1,0} slice(f32[4096,6]{1,0} %add.726), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_0.584 = f32[4096,3]{1,0} parameter(0)
  %slice.327 = f32[4096,2]{1,0} slice(f32[4096,3]{1,0} %param_0.584), slice={[0:4096], [0:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 2) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.326 = f32[4096,2]{1,0} slice(f32[4096,6]{1,0} %add.726), slice={[0:4096:1], [2:6:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 6) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.453 = f32[4096,2]{1,0} add(f32[4096,2]{1,0} %slice.327, f32[4096,2]{1,0} %slice.326), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.124 = f32[4096,3]{1,0} concatenate(f32[4096,1]{1,0} %slice.328, f32[4096,2]{1,0} %add.453), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_644 = f32[] constant(0)
  ROOT %pad.122 = f32[4096,6]{1,0} pad(f32[4096,3]{1,0} %concatenate.124, f32[] %constant_644), padding=0_0_0x0_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.106 (param_0.587: f32[4096,13]) -> f32[4096,3] {
  %param_0.587 = f32[4096,13]{1,0} parameter(0)
  %slice.392 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_0.587), slice={[0:4096:1], [0:12:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 12) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.391 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_0.587), slice={[0:4096:1], [1:13:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 13) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.729 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %slice.392, f32[4096,6]{1,0} %slice.391), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.334 = f32[4096,3]{1,0} slice(f32[4096,6]{1,0} %add.729), slice={[0:4096:1], [0:5:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 5) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.333 = f32[4096,3]{1,0} slice(f32[4096,6]{1,0} %add.729), slice={[0:4096:1], [1:6:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 6) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.457 = f32[4096,3]{1,0} add(f32[4096,3]{1,0} %slice.334, f32[4096,3]{1,0} %slice.333), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.332 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.457), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.331 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.457), slice={[0:4096:1], [0:2:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 2) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.330 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.457), slice={[0:4096:1], [1:3:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 3) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.456 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %slice.331, f32[4096,1]{1,0} %slice.330), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.329 = f32[4096,1]{1,0} slice(f32[4096,3]{1,0} %add.457), slice={[0:4096:1], [2:3:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 3) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.455 = f32[4096,1]{1,0} add(f32[4096,1]{1,0} %add.456, f32[4096,1]{1,0} %slice.329), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.125 = f32[4096,2]{1,0} concatenate(f32[4096,1]{1,0} %slice.332, f32[4096,1]{1,0} %add.455), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_645 = f32[] constant(0)
  %pad.124 = f32[4096,3]{1,0} pad(f32[4096,2]{1,0} %concatenate.125, f32[] %constant_645), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.123 = f32[4096,3]{1,0} pad(f32[4096,1]{1,0} %add.456, f32[] %constant_645), padding=0_0x1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %add.454 = f32[4096,3]{1,0} add(f32[4096,3]{1,0} %pad.124, f32[4096,3]{1,0} %pad.123), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.108 (param_0.582: f32[4096,55]) -> f32[4096,13] {
  %param_0.582 = f32[4096,55]{1,0} parameter(0)
  %slice.384 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_0.582), slice={[0:4096:1], [0:54:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 54) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.383 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_0.582), slice={[0:4096:1], [1:55:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 55) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.724 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %slice.384, f32[4096,27]{1,0} %slice.383), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.338 = f32[4096,13]{1,0} slice(f32[4096,27]{1,0} %add.724), slice={[0:4096:1], [0:26:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 26) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.337 = f32[4096,13]{1,0} slice(f32[4096,27]{1,0} %add.724), slice={[0:4096:1], [1:27:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 27) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %add.459 = f32[4096,13]{1,0} add(f32[4096,13]{1,0} %slice.338, f32[4096,13]{1,0} %slice.337), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.110 (param_0.578: pred[4096], param_1.793: f32[4096], param_2.763: f32[4096], param_3.627: f32[4096,221]) -> f32[4096,55] {
  %param_3.627 = f32[4096,221]{1,0} parameter(3)
  %negate.39 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_3.627), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_0.578 = pred[4096]{0} parameter(0)
  %param_2.763 = f32[4096]{0} parameter(2)
  %param_1.793 = f32[4096]{0} parameter(1)
  %constant_1256 = f32[] constant(0.001)
  %broadcast.1705 = f32[4096]{0} broadcast(f32[] %constant_1256), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.720 = f32[4096]{0} add(f32[4096]{0} %param_1.793, f32[4096]{0} %broadcast.1705), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.89 = f32[4096]{0} maximum(f32[4096]{0} %param_2.763, f32[4096]{0} %add.720), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.292 = f32[4096]{0} select(pred[4096]{0} %param_0.578, f32[4096]{0} %maximum.89, f32[4096]{0} %broadcast.1705), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %constant_1255 = f32[] constant(0)
  %broadcast.1704 = f32[4096]{0} broadcast(f32[] %constant_1255), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.291 = f32[4096]{0} select(pred[4096]{0} %param_0.578, f32[4096]{0} %param_1.793, f32[4096]{0} %broadcast.1704), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.252 = f32[4096]{0} subtract(f32[4096]{0} %select.292, f32[4096]{0} %select.291), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1254 = f32[] constant(0.00452488707)
  %broadcast.1703 = f32[4096]{0} broadcast(f32[] %constant_1254), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.516 = f32[4096]{0} multiply(f32[4096]{0} %subtract.252, f32[4096]{0} %broadcast.1703), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.1702 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.516), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.515 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.39, f32[4096,221]{1,0} %broadcast.1702), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %slice.376 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %multiply.515), slice={[0:4096:1], [0:220:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 220) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.375 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %multiply.515), slice={[0:4096:1], [1:221:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.719 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %slice.376, f32[4096,110]{1,0} %slice.375), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.342 = f32[4096,55]{1,0} slice(f32[4096,110]{1,0} %add.719), slice={[0:4096:1], [0:109:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 109) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.341 = f32[4096,55]{1,0} slice(f32[4096,110]{1,0} %add.719), slice={[0:4096:1], [1:110:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 110) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %add.461 = f32[4096,55]{1,0} add(f32[4096,55]{1,0} %slice.342, f32[4096,55]{1,0} %slice.341), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
}

%fused_computation.114 (param_0.270: f32[4096,221]) -> f32[4096,221] {
  %param_0.270 = f32[4096,221]{1,0} parameter(0)
  %compare.143 = pred[4096,221]{1,0} compare(f32[4096,221]{1,0} %param_0.270, f32[4096,221]{1,0} %param_0.270), direction=NE, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/ne" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %constant_647 = f32[] constant(0)
  %broadcast.1068 = f32[4096,221]{1,0} broadcast(f32[] %constant_647), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %maximum.61 = f32[4096,221]{1,0} maximum(f32[4096,221]{1,0} %param_0.270, f32[4096,221]{1,0} %broadcast.1068), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %abs.3 = f32[4096,221]{1,0} abs(f32[4096,221]{1,0} %param_0.270), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/abs" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %negate.22 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %abs.3), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %exponential.14 = f32[4096,221]{1,0} exponential(f32[4096,221]{1,0} %negate.22), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/exp" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %log-plus-one.3 = f32[4096,221]{1,0} log-plus-one(f32[4096,221]{1,0} %exponential.14), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/log1p" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %add.463 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %maximum.61, f32[4096,221]{1,0} %log-plus-one.3), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  ROOT %select.137 = f32[4096,221]{1,0} select(pred[4096,221]{1,0} %compare.143, f32[4096,221]{1,0} %param_0.270, f32[4096,221]{1,0} %add.463), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
}

%region_4.510 (Arg_0.511: f32[], Arg_1.512: f32[]) -> f32[] {
  %Arg_0.511 = f32[] parameter(0)
  %Arg_1.512 = f32[] parameter(1)
  ROOT %add.513 = f32[] add(f32[] %Arg_0.511, f32[] %Arg_1.512), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
}

%fused_computation.115 (param_0.558: f32[3,16,128], param_1.756: f32[3,4096,221], param_2.721: f32[3,16,128,128], param_3.584: f32[3,4096,221], param_4.358: f32[3,4096,221]) -> f32[4096,221] {
  %constant_820 = f32[] constant(1)
  %broadcast.1262 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_820), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.756 = f32[3,4096,221]{2,1,0} parameter(1)
  %floor.79 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_1.756), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.136 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_1.756, f32[3,4096,221]{2,1,0} %floor.79), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.135 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1262, f32[3,4096,221]{2,1,0} %subtract.136), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1261 = f32[3,4096,16,221]{2,1,0,3} broadcast(f32[3,4096,221]{2,1,0} %subtract.135), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_0.558 = f32[3,16,128]{2,1,0} parameter(0)
  %iota.119 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_816 = s32[] constant(0)
  %broadcast.1260 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_816), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/max" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.50 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.79), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_814 = s32[] constant(127)
  %broadcast.1259 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_814), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.39 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1260, s32[3,4096,221]{2,1,0} %convert.50, s32[3,4096,221]{2,1,0} %broadcast.1259), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.165 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.39, s32[3,4096,221]{2,1,0} %broadcast.1260), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_813 = s32[] constant(128)
  %broadcast.1258 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_813), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.544 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.39, s32[3,4096,221]{2,1,0} %broadcast.1258), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.177 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.165, s32[3,4096,221]{2,1,0} %add.544, s32[3,4096,221]{2,1,0} %clamp.39), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.172 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.177), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.158 = s32[3,4096,221,2]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.119, s32[3,4096,221,1]{2,1,0,3} %bitcast.172), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.87 = f32[3,4096,16,221]{2,1,0,3} gather(f32[3,16,128]{2,1,0} %param_0.558, s32[3,4096,221,2]{2,1,0,3} %concatenate.158), offset_dims={2}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,16,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 16, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.431 = f32[3,4096,16,221]{2,1,0,3} multiply(f32[3,4096,16,221]{2,1,0,3} %broadcast.1261, f32[3,4096,16,221]{2,1,0,3} %gather.87), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %copy.35 = f32[3,4096,16,221]{3,2,1,0} copy(f32[3,4096,16,221]{2,1,0,3} %multiply.431), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1257 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %subtract.136), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.118 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_815 = s32[] constant(1)
  %broadcast.1255 = s32[3,4096,221]{2,1,0} broadcast(s32[] %constant_815), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.543 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.50, s32[3,4096,221]{2,1,0} %broadcast.1255), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.38 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1260, s32[3,4096,221]{2,1,0} %add.543, s32[3,4096,221]{2,1,0} %broadcast.1259), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.164 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.38, s32[3,4096,221]{2,1,0} %broadcast.1260), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.542 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.38, s32[3,4096,221]{2,1,0} %broadcast.1258), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.176 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.164, s32[3,4096,221]{2,1,0} %add.542, s32[3,4096,221]{2,1,0} %clamp.38), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.171 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.176), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.157 = s32[3,4096,221,2]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.118, s32[3,4096,221,1]{3,2,1,0} %bitcast.171), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.86 = f32[3,4096,16,221]{3,2,1,0} gather(f32[3,16,128]{2,1,0} %param_0.558, s32[3,4096,221,2]{3,2,1,0} %concatenate.157), offset_dims={2}, collapsed_slice_dims={0,2}, start_index_map={0,2}, index_vector_dim=3, slice_sizes={1,16,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2), start_index_map=(0, 2)) slice_sizes=(1, 16, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.429 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1257, f32[3,4096,16,221]{3,2,1,0} %gather.86), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.541 = f32[3,4096,16,221]{3,2,1,0} add(f32[3,4096,16,221]{3,2,1,0} %copy.35, f32[3,4096,16,221]{3,2,1,0} %multiply.429), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_4.358 = f32[3,4096,221]{2,1,0} parameter(4)
  %floor.223 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_4.358), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.232 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_4.358, f32[3,4096,221]{2,1,0} %floor.223), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.231 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1262, f32[3,4096,221]{2,1,0} %subtract.232), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_3.584 = f32[3,4096,221]{2,1,0} parameter(3)
  %floor.222 = f32[3,4096,221]{2,1,0} floor(f32[3,4096,221]{2,1,0} %param_3.584), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/floor" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.230 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %param_3.584, f32[3,4096,221]{2,1,0} %floor.222), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %subtract.229 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %broadcast.1262, f32[3,4096,221]{2,1,0} %subtract.230), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/sub" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.481 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.231, f32[3,4096,221]{2,1,0} %subtract.229), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1622 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.481), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_2.721 = f32[3,16,128,128]{3,2,1,0} parameter(2)
  %iota.151 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.130 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.223), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.119 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1260, s32[3,4096,221]{2,1,0} %convert.130, s32[3,4096,221]{2,1,0} %broadcast.1259), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.245 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.119, s32[3,4096,221]{2,1,0} %broadcast.1260), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.697 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.119, s32[3,4096,221]{2,1,0} %broadcast.1258), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.259 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.245, s32[3,4096,221]{2,1,0} %add.697, s32[3,4096,221]{2,1,0} %clamp.119), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.244 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.259), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.129 = s32[3,4096,221]{2,1,0} convert(f32[3,4096,221]{2,1,0} %floor.222), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/convert_element_type[new_dtype=int32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.118 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1260, s32[3,4096,221]{2,1,0} %convert.129, s32[3,4096,221]{2,1,0} %broadcast.1259), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.244 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.118, s32[3,4096,221]{2,1,0} %broadcast.1260), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.696 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.118, s32[3,4096,221]{2,1,0} %broadcast.1258), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.258 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.244, s32[3,4096,221]{2,1,0} %add.696, s32[3,4096,221]{2,1,0} %clamp.118), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.243 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.258), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.190 = s32[3,4096,221,3]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.151, s32[3,4096,221,1]{2,1,0,3} %bitcast.244, s32[3,4096,221,1]{2,1,0,3} %bitcast.243), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.103 = f32[3,4096,16,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_2.721, s32[3,4096,221,3]{2,1,0,3} %concatenate.190), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.480 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1622, f32[3,4096,16,221]{3,2,1,0} %gather.103), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.479 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.231, f32[3,4096,221]{2,1,0} %subtract.230), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1614 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.479), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.150 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.242 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.259), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.694 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.129, s32[3,4096,221]{2,1,0} %broadcast.1255), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.116 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1260, s32[3,4096,221]{2,1,0} %add.694, s32[3,4096,221]{2,1,0} %broadcast.1259), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.242 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.116, s32[3,4096,221]{2,1,0} %broadcast.1260), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.693 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.116, s32[3,4096,221]{2,1,0} %broadcast.1258), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.256 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.242, s32[3,4096,221]{2,1,0} %add.693, s32[3,4096,221]{2,1,0} %clamp.116), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.241 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.256), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.189 = s32[3,4096,221,3]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.150, s32[3,4096,221,1]{3,2,1,0} %bitcast.242, s32[3,4096,221,1]{3,2,1,0} %bitcast.241), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.102 = f32[3,4096,16,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_2.721, s32[3,4096,221,3]{3,2,1,0} %concatenate.189), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.478 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1614, f32[3,4096,16,221]{3,2,1,0} %gather.102), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.692 = f32[3,4096,16,221]{3,2,1,0} add(f32[3,4096,16,221]{3,2,1,0} %multiply.480, f32[3,4096,16,221]{3,2,1,0} %multiply.478), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.477 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.232, f32[3,4096,221]{2,1,0} %subtract.229), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1605 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.477), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.149 = s32[3,4096,221,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.691 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %convert.130, s32[3,4096,221]{2,1,0} %broadcast.1255), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %clamp.115 = s32[3,4096,221]{2,1,0} clamp(s32[3,4096,221]{2,1,0} %broadcast.1260, s32[3,4096,221]{2,1,0} %add.691, s32[3,4096,221]{2,1,0} %broadcast.1259), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(vmap(vmap(clip)))/jit(jit_vmap(vmap(clip)))/min" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %compare.241 = pred[3,4096,221]{2,1,0} compare(s32[3,4096,221]{2,1,0} %clamp.115, s32[3,4096,221]{2,1,0} %broadcast.1260), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/lt" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.690 = s32[3,4096,221]{2,1,0} add(s32[3,4096,221]{2,1,0} %clamp.115, s32[3,4096,221]{2,1,0} %broadcast.1258), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %select.255 = s32[3,4096,221]{2,1,0} select(pred[3,4096,221]{2,1,0} %compare.241, s32[3,4096,221]{2,1,0} %add.690, s32[3,4096,221]{2,1,0} %clamp.115), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/select_n" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.240 = s32[3,4096,221,1]{2,1,0,3} bitcast(s32[3,4096,221]{2,1,0} %select.255), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.188 = s32[3,4096,221,3]{2,1,0,3} concatenate(s32[3,4096,221,1]{2,1,0,3} %iota.149, s32[3,4096,221,1]{2,1,0,3} %bitcast.240, s32[3,4096,221,1]{2,1,0,3} %bitcast.243), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.101 = f32[3,4096,16,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_2.721, s32[3,4096,221,3]{2,1,0,3} %concatenate.188), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.476 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1605, f32[3,4096,16,221]{3,2,1,0} %gather.101), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.688 = f32[3,4096,16,221]{3,2,1,0} add(f32[3,4096,16,221]{3,2,1,0} %add.692, f32[3,4096,16,221]{3,2,1,0} %multiply.476), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.475 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %subtract.232, f32[3,4096,221]{2,1,0} %subtract.230), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %broadcast.1597 = f32[3,4096,16,221]{3,2,1,0} broadcast(f32[3,4096,221]{2,1,0} %multiply.475), dimensions={0,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %iota.148 = s32[3,4096,221,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/iota[dtype=int32 shape=(3, 4096, 221, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.238 = s32[3,4096,221,1]{3,2,1,0} bitcast(s32[3,4096,221]{2,1,0} %select.255), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/broadcast_in_dim[shape=(3, 4096, 221, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.187 = s32[3,4096,221,3]{3,2,1,0} concatenate(s32[3,4096,221,1]{3,2,1,0} %iota.148, s32[3,4096,221,1]{3,2,1,0} %bitcast.238, s32[3,4096,221,1]{3,2,1,0} %bitcast.241), dimensions={3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/concatenate[dimension=3]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %gather.100 = f32[3,4096,16,221]{3,2,1,0} gather(f32[3,16,128,128]{3,2,1,0} %param_2.721, s32[3,4096,221,3]{3,2,1,0} %concatenate.187), offset_dims={2}, collapsed_slice_dims={0,2,3}, start_index_map={0,2,3}, index_vector_dim=3, slice_sizes={1,16,1,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(2,), collapsed_slice_dims=(0, 2, 3), start_index_map=(0, 2, 3)) slice_sizes=(1, 16, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.474 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %broadcast.1597, f32[3,4096,16,221]{3,2,1,0} %gather.100), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.683 = f32[3,4096,16,221]{3,2,1,0} add(f32[3,4096,16,221]{3,2,1,0} %add.688, f32[3,4096,16,221]{3,2,1,0} %multiply.474), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %multiply.378 = f32[3,4096,16,221]{3,2,1,0} multiply(f32[3,4096,16,221]{3,2,1,0} %add.541, f32[3,4096,16,221]{3,2,1,0} %add.683), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %constant_648 = f32[] constant(0)
  %reduce.46 = f32[4096,221]{1,0} reduce(f32[3,4096,16,221]{3,2,1,0} %multiply.378, f32[] %constant_648), dimensions={0,2}, to_apply=%region_4.510, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %constant_464 = f32[] constant(10)
  %broadcast.925 = f32[4096,221]{1,0} broadcast(f32[] %constant_464), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  ROOT %add.464 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %reduce.46, f32[4096,221]{1,0} %broadcast.925), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
}

%fused_computation.130 (param_0.507: f32[3,4096,221], param_1.671: s32[3,1], param_2.614: s32[3,1]) -> (f32[3,4096,221], f32[3,4096,221], f32[3,4096,221]) {
  %param_0.507 = f32[3,4096,221]{2,1,0} parameter(0)
  %bitcast.156 = f32[1,3,4096,221]{3,2,0,1} bitcast(f32[3,4096,221]{2,1,0} %param_0.507), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 221) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %param_2.614 = s32[3,1]{1,0} parameter(2)
  %gather.79 = f32[3,4096,221]{2,1,0} gather(f32[3,4096,221]{2,1,0} %param_0.507, s32[3,1]{1,0} %param_2.614), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,4096,221}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 4096, 221) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=58}
  %bitcast.155 = f32[1,3,4096,221]{3,2,0,1} bitcast(f32[3,4096,221]{2,1,0} %gather.79), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 221) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %param_1.671 = s32[3,1]{1,0} parameter(1)
  %gather.78 = f32[3,4096,221]{2,1,0} gather(f32[3,4096,221]{2,1,0} %param_0.507, s32[3,1]{1,0} %param_1.671), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,4096,221}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 4096, 221) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=59}
  %bitcast.154 = f32[1,3,4096,221]{3,2,0,1} bitcast(f32[3,4096,221]{2,1,0} %gather.78), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 221) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %concatenate.142 = f32[3,3,4096,221]{3,2,0,1} concatenate(f32[1,3,4096,221]{3,2,0,1} %bitcast.156, f32[1,3,4096,221]{3,2,0,1} %bitcast.155, f32[1,3,4096,221]{3,2,0,1} %bitcast.154), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %constant_742 = f32[] constant(1)
  %broadcast.1179 = f32[3,3,4096,221]{3,2,0,1} broadcast(f32[] %constant_742), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=159}
  %add.514 = f32[3,3,4096,221]{3,2,0,1} add(f32[3,3,4096,221]{3,2,0,1} %concatenate.142, f32[3,3,4096,221]{3,2,0,1} %broadcast.1179), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=159}
  %constant_741 = f32[] constant(63.5)
  %broadcast.1178 = f32[3,3,4096,221]{3,2,0,1} broadcast(f32[] %constant_741), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=162}
  %multiply.421 = f32[3,3,4096,221]{3,2,0,1} multiply(f32[3,3,4096,221]{3,2,0,1} %add.514, f32[3,3,4096,221]{3,2,0,1} %broadcast.1178), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=162}
  %transpose.39 = f32[3,4096,3,221]{3,1,0,2} transpose(f32[3,3,4096,221]{3,2,0,1} %multiply.421), dimensions={0,2,1,3}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/transpose[permutation=(0, 2, 1, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=165}
  %slice.346 = f32[3,4096,1,221]{3,1,0,2} slice(f32[3,4096,3,221]{3,1,0,2} %transpose.39), slice={[0:3], [0:4096], [1:2], [0:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 0, 0) limit_indices=(3, 4096, 1, 221) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.118 = f32[3,4096,221]{2,1,0} bitcast(f32[3,4096,1,221]{3,1,0,2} %slice.346), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.347.clone.1 = f32[3,4096,1,221]{3,1,0,2} slice(f32[3,4096,3,221]{3,1,0,2} %transpose.39), slice={[0:3], [0:4096], [0:1], [0:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1, 2, 3), collapsed_slice_dims=(), start_index_map=(2,)) slice_sizes=(3, 4096, 1, 221) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=165}
  %bitcast.121.clone.1 = f32[3,4096,221]{2,1,0} bitcast(f32[3,4096,1,221]{3,1,0,2} %slice.347.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %slice.345.clone.1 = f32[3,4096,1,221]{3,1,0,2} slice(f32[3,4096,3,221]{3,1,0,2} %transpose.39), slice={[0:3], [0:4096], [2:3], [0:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/slice[start_indices=(0, 0, 1, 0) limit_indices=(3, 4096, 2, 221) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.117.clone.1 = f32[3,4096,221]{2,1,0} bitcast(f32[3,4096,1,221]{3,1,0,2} %slice.345.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  ROOT %tuple.80 = (f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}) tuple(f32[3,4096,221]{2,1,0} %bitcast.118, f32[3,4096,221]{2,1,0} %bitcast.121.clone.1, f32[3,4096,221]{2,1,0} %bitcast.117.clone.1)
}

%fused_computation.137 (param_0.433: u32[111], param_1.582: u32[111]) -> f32[221] {
  %param_0.433 = u32[111]{0} parameter(0)
  %param_1.582 = u32[111]{0} parameter(1)
  %concatenate.133 = u32[222]{0} concatenate(u32[111]{0} %param_0.433, u32[111]{0} %param_1.582), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %slice.348 = u32[221]{0} slice(u32[222]{0} %concatenate.133), slice={[0:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(), start_index_map=(0,)) slice_sizes=(221,) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_631 = u32[] constant(9)
  %broadcast.1065 = u32[221]{0} broadcast(u32[] %constant_631), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/shift_right_logical" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %shift-right-logical.2 = u32[221]{0} shift-right-logical(u32[221]{0} %slice.348, u32[221]{0} %broadcast.1065), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/shift_right_logical" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_632 = u32[] constant(1065353216)
  %broadcast.1066 = u32[221]{0} broadcast(u32[] %constant_632), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/or" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %or.2 = u32[221]{0} or(u32[221]{0} %shift-right-logical.2, u32[221]{0} %broadcast.1066), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/or" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %bitcast-convert.20 = f32[221]{0} bitcast-convert(u32[221]{0} %or.2), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/bitcast_convert_type[new_dtype=float32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_637 = f32[] constant(-1)
  %broadcast.1067 = f32[221]{0} broadcast(f32[] %constant_637), dimensions={}
  %add.481 = f32[221]{0} add(f32[221]{0} %bitcast-convert.20, f32[221]{0} %broadcast.1067), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %constant_466 = f32[] constant(1.17549435e-38)
  %broadcast.935 = f32[221]{0} broadcast(f32[] %constant_466), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %maximum.62 = f32[221]{0} maximum(f32[221]{0} %add.481, f32[221]{0} %broadcast.935), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %log.6 = f32[221]{0} log(f32[221]{0} %maximum.62), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/log" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %negate.23 = f32[221]{0} negate(f32[221]{0} %log.6), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  ROOT %log.5 = f32[221]{0} log(f32[221]{0} %negate.23), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/log" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
}

%fused_computation.139 (param_0.606: u32[2], param_1.825: u32[2]) -> (u32[111], u32[111], u32[111], u32[111]) {
  %param_0.606 = u32[2]{0} parameter(0)
  %param_1.825 = u32[2]{0} parameter(1)
  %concatenate.202 = u32[4]{0} concatenate(u32[2]{0} %param_0.606, u32[2]{0} %param_1.825), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %bitcast.248 = u32[2,2]{1,0} bitcast(u32[4]{0} %concatenate.202), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(2, 2) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %slice.351 = u32[1,1]{1,0} slice(u32[2,2]{1,0} %bitcast.248), slice={[1:2], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(0,) limit_indices=(1,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %bitcast.127 = u32[] bitcast(u32[1,1]{1,0} %slice.351), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %broadcast.937 = u32[111]{0} broadcast(u32[] %bitcast.127), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %slice.350.clone.1 = u32[1,2]{1,0} slice(u32[2,2]{1,0} %bitcast.248), slice={[1:2], [0:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(1, 0) limit_indices=(2, 2) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %bitcast.126.clone.1 = u32[2]{0} bitcast(u32[1,2]{1,0} %slice.350.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %slice.349.clone.1 = u32[1]{0} slice(u32[2]{0} %bitcast.126.clone.1), slice={[1:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(1,) limit_indices=(2,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %bitcast.125.clone.1 = u32[] bitcast(u32[1]{0} %slice.349.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %broadcast.936.clone.1 = u32[111]{0} broadcast(u32[] %bitcast.125.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %slice.356.clone.1 = u32[1]{0} slice(u32[2]{0} %param_0.606), slice={[1:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(1,) limit_indices=(2,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %bitcast.131.clone.1 = u32[] bitcast(u32[1]{0} %slice.356.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %broadcast.952.clone.1 = u32[111]{0} broadcast(u32[] %bitcast.131.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %slice.357.clone.1 = u32[1]{0} slice(u32[2]{0} %param_0.606), slice={[0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(0,) limit_indices=(1,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %bitcast.132.clone.1 = u32[] bitcast(u32[1]{0} %slice.357.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %broadcast.953.clone.1 = u32[111]{0} broadcast(u32[] %bitcast.132.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  ROOT %tuple.85 = (u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}) tuple(u32[111]{0} %broadcast.937, u32[111]{0} %broadcast.936.clone.1, u32[111]{0} %broadcast.952.clone.1, u32[111]{0} %broadcast.953.clone.1)
}

%fused_computation.141 (param_0.470: f32[2,3], param_1.641: u32[111], param_2.588: u32[111], param_3.449: f32[4096,3], param_4.248: f32[4096,3], param_5.147: pred[4096], param_6.93: f32[4096], param_7.42: f32[4096]) -> f32[3,4096,221] {
  %param_4.248 = f32[4096,3]{1,0} parameter(4)
  %broadcast.948 = f32[4096,3,221]{2,0,1} broadcast(f32[4096,3]{1,0} %param_4.248), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=335}
  %param_3.449 = f32[4096,3]{1,0} parameter(3)
  %broadcast.947 = f32[4096,3,221]{2,0,1} broadcast(f32[4096,3]{1,0} %param_3.449), dimensions={0,1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=335}
  %param_5.147 = pred[4096]{0} parameter(5)
  %param_6.93 = f32[4096]{0} parameter(6)
  %constant_679 = f32[] constant(0)
  %broadcast.1129 = f32[4096]{0} broadcast(f32[] %constant_679), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.147 = f32[4096]{0} select(pred[4096]{0} %param_5.147, f32[4096]{0} %param_6.93, f32[4096]{0} %broadcast.1129), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %broadcast.946 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %select.147), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=331}
  %iota.99 = s32[221]{0} iota(), iota_dimension=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/iota[dtype=int32 shape=(221,) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=322}
  %convert.20 = f32[221]{0} convert(s32[221]{0} %iota.99), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %param_1.641 = u32[111]{0} parameter(1)
  %param_2.588 = u32[111]{0} parameter(2)
  %concatenate.135 = u32[222]{0} concatenate(u32[111]{0} %param_1.641, u32[111]{0} %param_2.588), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %slice.354 = u32[221]{0} slice(u32[222]{0} %concatenate.135), slice={[0:221]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(), start_index_map=(0,)) slice_sizes=(221,) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %constant_676 = u32[] constant(9)
  %broadcast.1101 = u32[221]{0} broadcast(u32[] %constant_676), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/shift_right_logical" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %shift-right-logical.3 = u32[221]{0} shift-right-logical(u32[221]{0} %slice.354, u32[221]{0} %broadcast.1101), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/shift_right_logical" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %constant_677 = u32[] constant(1065353216)
  %broadcast.1102 = u32[221]{0} broadcast(u32[] %constant_677), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/or" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %or.3 = u32[221]{0} or(u32[221]{0} %shift-right-logical.3, u32[221]{0} %broadcast.1102), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/or" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %bitcast-convert.21 = f32[221]{0} bitcast-convert(u32[221]{0} %or.3), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/bitcast_convert_type[new_dtype=float32]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %constant_678 = f32[] constant(-1)
  %broadcast.1103 = f32[221]{0} broadcast(f32[] %constant_678), dimensions={}
  %add.487 = f32[221]{0} add(f32[221]{0} %bitcast-convert.21, f32[221]{0} %broadcast.1103), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %broadcast.945 = f32[221]{0} broadcast(f32[] %constant_679), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %maximum.63 = f32[221]{0} maximum(f32[221]{0} %add.487, f32[221]{0} %broadcast.945), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %add.486 = f32[221]{0} add(f32[221]{0} %convert.20, f32[221]{0} %maximum.63), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %broadcast.944 = f32[4096,221]{1,0} broadcast(f32[221]{0} %add.486), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %param_7.42 = f32[4096]{0} parameter(7)
  %constant_713 = f32[] constant(0.001)
  %broadcast.1146 = f32[4096]{0} broadcast(f32[] %constant_713), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.503 = f32[4096]{0} add(f32[4096]{0} %param_6.93, f32[4096]{0} %broadcast.1146), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.69 = f32[4096]{0} maximum(f32[4096]{0} %param_7.42, f32[4096]{0} %add.503), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.153 = f32[4096]{0} select(pred[4096]{0} %param_5.147, f32[4096]{0} %maximum.69, f32[4096]{0} %broadcast.1146), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %subtract.112 = f32[4096]{0} subtract(f32[4096]{0} %select.153, f32[4096]{0} %select.147), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_711 = f32[] constant(0.00452488707)
  %broadcast.1142 = f32[4096]{0} broadcast(f32[] %constant_711), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.410 = f32[4096]{0} multiply(f32[4096]{0} %subtract.112, f32[4096]{0} %broadcast.1142), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.1139 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.410), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.393 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %broadcast.944, f32[4096,221]{1,0} %broadcast.1139), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %add.485 = f32[4096,221]{1,0} add(f32[4096,221]{1,0} %broadcast.946, f32[4096,221]{1,0} %multiply.393), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=331}
  %broadcast.943 = f32[4096,3,221]{2,0,1} broadcast(f32[4096,221]{1,0} %add.485), dimensions={0,2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=335}
  %multiply.392 = f32[4096,3,221]{2,0,1} multiply(f32[4096,3,221]{2,0,1} %broadcast.947, f32[4096,3,221]{2,0,1} %broadcast.943), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=335}
  %add.484 = f32[4096,3,221]{2,0,1} add(f32[4096,3,221]{2,0,1} %broadcast.948, f32[4096,3,221]{2,0,1} %multiply.392), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=335}
  %transpose.33 = f32[3,4096,221]{2,1,0} transpose(f32[4096,3,221]{2,0,1} %add.484), dimensions={1,0,2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/transpose[permutation=(1, 0, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=114}
  %param_0.470 = f32[2,3]{1,0} parameter(0)
  %slice.353 = f32[1,3]{1,0} slice(f32[2,3]{1,0} %param_0.470), slice={[0:1], [0:3]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 3) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %bitcast.130 = f32[3]{0} bitcast(f32[1,3]{1,0} %slice.353), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %broadcast.942 = f32[3,4096,221]{2,1,0} broadcast(f32[3]{0} %bitcast.130), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %subtract.74 = f32[3,4096,221]{2,1,0} subtract(f32[3,4096,221]{2,1,0} %transpose.33, f32[3,4096,221]{2,1,0} %broadcast.942), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %slice.352 = f32[1,3]{1,0} slice(f32[2,3]{1,0} %param_0.470), slice={[1:2], [0:3]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(0,), start_index_map=(0,)) slice_sizes=(1, 3) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %subtract.73 = f32[1,3]{1,0} subtract(f32[1,3]{1,0} %slice.352, f32[1,3]{1,0} %slice.353), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %bitcast.129 = f32[3]{0} bitcast(f32[1,3]{1,0} %subtract.73), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %broadcast.941 = f32[3,4096,221]{2,1,0} broadcast(f32[3]{0} %bitcast.129), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %divide.48 = f32[3,4096,221]{2,1,0} divide(f32[3,4096,221]{2,1,0} %subtract.74, f32[3,4096,221]{2,1,0} %broadcast.941), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %constant_468 = f32[] constant(-0.5)
  %broadcast.940 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_468), dimensions={}
  %add.482 = f32[3,4096,221]{2,1,0} add(f32[3,4096,221]{2,1,0} %divide.48, f32[3,4096,221]{2,1,0} %broadcast.940), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=129}
  %constant_680 = f32[] constant(2)
  %broadcast.938 = f32[3,4096,221]{2,1,0} broadcast(f32[] %constant_680), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=128}
  ROOT %multiply.391 = f32[3,4096,221]{2,1,0} multiply(f32[3,4096,221]{2,1,0} %add.482, f32[3,4096,221]{2,1,0} %broadcast.938), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=128}
}

%fused_computation.149 (param_0.364: u32[2]) -> (u32[2], u32[2]) {
  %param_0.364 = u32[2]{0} parameter(0)
  %slice.363 = u32[1]{0} slice(u32[2]{0} %param_0.364), slice={[0:1]}, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(0,) limit_indices=(1,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %bitcast.136 = u32[] bitcast(u32[1]{0} %slice.363), metadata={op_name="jit(training_step)/jit(main)/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %broadcast.957 = u32[2]{0} broadcast(u32[] %bitcast.136), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %slice.362.clone.1 = u32[1]{0} slice(u32[2]{0} %param_0.364), slice={[1:2]}, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(1,) limit_indices=(2,) strides=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %bitcast.135.clone.1 = u32[] bitcast(u32[1]{0} %slice.362.clone.1), metadata={op_name="jit(training_step)/jit(main)/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %broadcast.956.clone.1 = u32[2]{0} broadcast(u32[] %bitcast.135.clone.1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  ROOT %tuple.88 = (u32[2]{0}, u32[2]{0}) tuple(u32[2]{0} %broadcast.957, u32[2]{0} %broadcast.956.clone.1)
}

%fused_computation.152 (param_0.401: f32[135168,128]) -> f32[135168,128] {
  %param_0.401 = f32[135168,128]{1,0} parameter(0)
  %constant_598 = f32[] constant(0)
  %broadcast.1019 = f32[135168,128]{1,0} broadcast(f32[] %constant_598), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gt" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  ROOT %maximum.66 = f32[135168,128]{1,0} maximum(f32[135168,128]{1,0} %param_0.401, f32[135168,128]{1,0} %broadcast.1019), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(relu)/jit(jit_relu)/max" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=105}
}

%fused_computation.153 (param_0.402: f32[135168,128]) -> f32[135168,128] {
  %param_0.402 = f32[135168,128]{1,0} parameter(0)
  %constant_599 = f32[] constant(0)
  %broadcast.1020 = f32[135168,128]{1,0} broadcast(f32[] %constant_599), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/gt" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  ROOT %maximum.67 = f32[135168,128]{1,0} maximum(f32[135168,128]{1,0} %param_0.402, f32[135168,128]{1,0} %broadcast.1020), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(relu)/jit(jit_relu)/max" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
}

%fused_computation.159 (param_0.754: f32[4096,3], param_1.995: f32[4096,3], param_2.1045: f32[2,3]) -> (f32[4096,3], f32[4096,3]) {
  %param_2.1045 = f32[2,3]{1,0} parameter(2)
  %broadcast.1124 = f32[4096,2,3]{2,1,0} broadcast(f32[2,3]{1,0} %param_2.1045), dimensions={1,2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=361}
  %param_1.995 = f32[4096,3]{1,0} parameter(1)
  %broadcast.1122 = f32[4096,2,3]{2,1,0} broadcast(f32[4096,3]{1,0} %param_1.995), dimensions={0,2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=361}
  %subtract.110 = f32[4096,2,3]{2,1,0} subtract(f32[4096,2,3]{2,1,0} %broadcast.1124, f32[4096,2,3]{2,1,0} %broadcast.1122), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=361}
  %param_0.754 = f32[4096,3]{1,0} parameter(0)
  %constant_700 = f32[] constant(1e-08)
  %broadcast.1121 = f32[4096,3]{1,0} broadcast(f32[] %constant_700), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=363}
  %add.501 = f32[4096,3]{1,0} add(f32[4096,3]{1,0} %param_0.754, f32[4096,3]{1,0} %broadcast.1121), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=363}
  %broadcast.1120 = f32[4096,2,3]{2,1,0} broadcast(f32[4096,3]{1,0} %add.501), dimensions={0,2}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=362}
  %divide.53 = f32[4096,2,3]{2,1,0} divide(f32[4096,2,3]{2,1,0} %subtract.110, f32[4096,2,3]{2,1,0} %broadcast.1120), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=362}
  %constant_1811 = f32[] constant(-inf)
  %reduce.49 = f32[4096,3]{1,0} reduce(f32[4096,2,3]{2,1,0} %divide.53, f32[] %constant_1811), dimensions={1}, to_apply=%region_1.225, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_max[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=368}
  %constant_1819 = f32[] constant(inf)
  %reduce.48.clone.1 = f32[4096,3]{1,0} reduce(f32[4096,2,3]{2,1,0} %divide.53, f32[] %constant_1819), dimensions={1}, to_apply=%region_0.221, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_min[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=367}
  ROOT %tuple.83 = (f32[4096,3]{1,0}, f32[4096,3]{1,0}) tuple(f32[4096,3]{1,0} %reduce.49, f32[4096,3]{1,0} %reduce.48.clone.1)
}

%fused_computation.161 (param_0.759: f32[4096,13], param_1.999: f32[4096,6], param_2.1049: f32[4096,6], param_3.925: f32[4096,55]) -> (f32[4096,27], f32[4096,27]) {
  %param_0.759 = f32[4096,13]{1,0} parameter(0)
  %slice.400 = f32[4096,1]{1,0} slice(f32[4096,13]{1,0} %param_0.759), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_1.999 = f32[4096,6]{1,0} parameter(1)
  %param_2.1049 = f32[4096,6]{1,0} parameter(2)
  %add.741 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %param_1.999, f32[4096,6]{1,0} %param_2.1049), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.399 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_0.759), slice={[0:4096:1], [2:13:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 13) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.740 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %add.741, f32[4096,6]{1,0} %slice.399), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.194 = f32[4096,7]{1,0} concatenate(f32[4096,1]{1,0} %slice.400, f32[4096,6]{1,0} %add.740), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_1266 = f32[] constant(0)
  %pad.133 = f32[4096,13]{1,0} pad(f32[4096,7]{1,0} %concatenate.194, f32[] %constant_1266), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.132 = f32[4096,13]{1,0} pad(f32[4096,6]{1,0} %add.741, f32[] %constant_1266), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.739 = f32[4096,13]{1,0} add(f32[4096,13]{1,0} %pad.133, f32[4096,13]{1,0} %pad.132), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.129 = f32[4096,27]{1,0} pad(f32[4096,13]{1,0} %add.739, f32[] %constant_1266), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_3.925 = f32[4096,55]{1,0} parameter(3)
  %slice.380.clone.1 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_3.925), slice={[0:4096:1], [0:54:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 54) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.379.clone.1 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_3.925), slice={[0:4096:1], [1:55:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 55) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.722.clone.1 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %slice.380.clone.1, f32[4096,27]{1,0} %slice.379.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.323.clone.1 = f32[4096,1]{1,0} slice(f32[4096,27]{1,0} %add.722.clone.1), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.322.clone.1 = f32[4096,13]{1,0} slice(f32[4096,27]{1,0} %add.722.clone.1), slice={[0:4096:1], [2:27:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 27) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.446.clone.1 = f32[4096,13]{1,0} add(f32[4096,13]{1,0} %add.739, f32[4096,13]{1,0} %slice.322.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.122.clone.1 = f32[4096,14]{1,0} concatenate(f32[4096,1]{1,0} %slice.323.clone.1, f32[4096,13]{1,0} %add.446.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.119.clone.1 = f32[4096,27]{1,0} pad(f32[4096,14]{1,0} %concatenate.122.clone.1, f32[] %constant_1266), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %tuple.78 = (f32[4096,27]{1,0}, f32[4096,27]{1,0}) tuple(f32[4096,27]{1,0} %pad.129, f32[4096,27]{1,0} %pad.119.clone.1)
}

%fused_computation.162 (param_0.758: f32[4096,55], param_1.998: f32[4096,27], param_2.1048: f32[4096,27], param_3.924: pred[4096], param_4.696: f32[4096], param_5.651: f32[4096], param_6.619: f32[4096,221]) -> (f32[4096,110], f32[4096,110]) {
  %param_0.758 = f32[4096,55]{1,0} parameter(0)
  %slice.408 = f32[4096,1]{1,0} slice(f32[4096,55]{1,0} %param_0.758), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_1.998 = f32[4096,27]{1,0} parameter(1)
  %param_2.1048 = f32[4096,27]{1,0} parameter(2)
  %add.753 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %param_1.998, f32[4096,27]{1,0} %param_2.1048), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.407 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_0.758), slice={[0:4096:1], [2:55:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 55) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.752 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %add.753, f32[4096,27]{1,0} %slice.407), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.198 = f32[4096,28]{1,0} concatenate(f32[4096,1]{1,0} %slice.408, f32[4096,27]{1,0} %add.752), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_1272 = f32[] constant(0)
  %pad.142 = f32[4096,55]{1,0} pad(f32[4096,28]{1,0} %concatenate.198, f32[] %constant_1272), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.141 = f32[4096,55]{1,0} pad(f32[4096,27]{1,0} %add.753, f32[] %constant_1272), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.751 = f32[4096,55]{1,0} add(f32[4096,55]{1,0} %pad.142, f32[4096,55]{1,0} %pad.141), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.138 = f32[4096,110]{1,0} pad(f32[4096,55]{1,0} %add.751, f32[] %constant_1272), padding=0_0_0x1_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_6.619 = f32[4096,221]{1,0} parameter(6)
  %negate.37.clone.1 = f32[4096,221]{1,0} negate(f32[4096,221]{1,0} %param_6.619), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/neg" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %param_3.924 = pred[4096]{0} parameter(3)
  %param_5.651 = f32[4096]{0} parameter(5)
  %param_4.696 = f32[4096]{0} parameter(4)
  %constant_1249_clone_1 = f32[] constant(0.001)
  %broadcast.1697.clone.1 = f32[4096]{0} broadcast(f32[] %constant_1249_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %add.715.clone.1 = f32[4096]{0} add(f32[4096]{0} %param_4.696, f32[4096]{0} %broadcast.1697.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %maximum.87.clone.1 = f32[4096]{0} maximum(f32[4096]{0} %param_5.651, f32[4096]{0} %add.715.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=374}
  %select.287.clone.1 = f32[4096]{0} select(pred[4096]{0} %param_3.924, f32[4096]{0} %maximum.87.clone.1, f32[4096]{0} %broadcast.1697.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=379}
  %broadcast.1696.clone.1 = f32[4096]{0} broadcast(f32[] %constant_1272), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %select.286.clone.1 = f32[4096]{0} select(pred[4096]{0} %param_3.924, f32[4096]{0} %param_4.696, f32[4096]{0} %broadcast.1696.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(vmap(_where))/jit(jit_vmap(_where))/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=378}
  %subtract.250.clone.1 = f32[4096]{0} subtract(f32[4096]{0} %select.287.clone.1, f32[4096]{0} %select.286.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %constant_1247_clone_1 = f32[] constant(0.00452488707)
  %broadcast.1695.clone.1 = f32[4096]{0} broadcast(f32[] %constant_1247_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %multiply.512.clone.1 = f32[4096]{0} multiply(f32[4096]{0} %subtract.250.clone.1, f32[4096]{0} %broadcast.1695.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/div" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=319}
  %broadcast.1694.clone.1 = f32[4096,221]{1,0} broadcast(f32[4096]{0} %multiply.512.clone.1), dimensions={0}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=330}
  %multiply.511.clone.1 = f32[4096,221]{1,0} multiply(f32[4096,221]{1,0} %negate.37.clone.1, f32[4096,221]{1,0} %broadcast.1694.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=271}
  %slice.372.clone.1 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %multiply.511.clone.1), slice={[0:4096:1], [0:220:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 220) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.371.clone.1 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %multiply.511.clone.1), slice={[0:4096:1], [1:221:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.714.clone.1 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %slice.372.clone.1, f32[4096,110]{1,0} %slice.371.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.319.clone.1 = f32[4096,1]{1,0} slice(f32[4096,110]{1,0} %add.714.clone.1), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.318.clone.1 = f32[4096,54]{1,0} slice(f32[4096,55]{1,0} %add.751), slice={[0:4096], [0:54]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 0) limit_indices=(4096, 54) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.317.clone.1 = f32[4096,54]{1,0} slice(f32[4096,110]{1,0} %add.714.clone.1), slice={[0:4096:1], [2:110:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/slice[start_indices=(0, 2) limit_indices=(4096, 110) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.442.clone.1 = f32[4096,54]{1,0} add(f32[4096,54]{1,0} %slice.318.clone.1, f32[4096,54]{1,0} %slice.317.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.120.clone.1 = f32[4096,55]{1,0} concatenate(f32[4096,1]{1,0} %slice.319.clone.1, f32[4096,54]{1,0} %add.442.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.116.clone.1 = f32[4096,110]{1,0} pad(f32[4096,55]{1,0} %concatenate.120.clone.1, f32[] %constant_1272), padding=0_0_0x0_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %tuple.77 = (f32[4096,110]{1,0}, f32[4096,110]{1,0}) tuple(f32[4096,110]{1,0} %pad.138, f32[4096,110]{1,0} %pad.116.clone.1)
}

%fused_computation.164 (param_0.757: f32[4096,13], param_1.997: f32[4096,6], param_2.1047: f32[4096,6], param_3.923: f32[4096,55]) -> (f32[4096,27], f32[4096,27]) {
  %param_0.757 = f32[4096,13]{1,0} parameter(0)
  %slice.686 = f32[4096,1]{1,0} slice(f32[4096,13]{1,0} %param_0.757), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_1.997 = f32[4096,6]{1,0} parameter(1)
  %param_2.1047 = f32[4096,6]{1,0} parameter(2)
  %add.979 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %param_1.997, f32[4096,6]{1,0} %param_2.1047), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.685 = f32[4096,6]{1,0} slice(f32[4096,13]{1,0} %param_0.757), slice={[0:4096:1], [2:13:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 13) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.978 = f32[4096,6]{1,0} add(f32[4096,6]{1,0} %add.979, f32[4096,6]{1,0} %slice.685), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.255 = f32[4096,7]{1,0} concatenate(f32[4096,1]{1,0} %slice.686, f32[4096,6]{1,0} %add.978), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_1764 = f32[] constant(0)
  %pad.155 = f32[4096,13]{1,0} pad(f32[4096,7]{1,0} %concatenate.255, f32[] %constant_1764), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.154 = f32[4096,13]{1,0} pad(f32[4096,6]{1,0} %add.979, f32[] %constant_1764), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.977 = f32[4096,13]{1,0} add(f32[4096,13]{1,0} %pad.155, f32[4096,13]{1,0} %pad.154), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.151 = f32[4096,27]{1,0} pad(f32[4096,13]{1,0} %add.977, f32[] %constant_1764), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_3.923 = f32[4096,55]{1,0} parameter(3)
  %slice.666.clone.1 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_3.923), slice={[0:4096:1], [0:54:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 54) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.665.clone.1 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_3.923), slice={[0:4096:1], [1:55:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 55) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.961.clone.1 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %slice.666.clone.1, f32[4096,27]{1,0} %slice.665.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.283.clone.1 = f32[4096,1]{1,0} slice(f32[4096,27]{1,0} %add.961.clone.1), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.282.clone.1 = f32[4096,13]{1,0} slice(f32[4096,27]{1,0} %add.961.clone.1), slice={[0:4096:1], [2:27:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 27) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.350.clone.1 = f32[4096,13]{1,0} add(f32[4096,13]{1,0} %add.977, f32[4096,13]{1,0} %slice.282.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.104.clone.1 = f32[4096,14]{1,0} concatenate(f32[4096,1]{1,0} %slice.283.clone.1, f32[4096,13]{1,0} %add.350.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.101.clone.1 = f32[4096,27]{1,0} pad(f32[4096,14]{1,0} %concatenate.104.clone.1, f32[] %constant_1764), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %tuple.50 = (f32[4096,27]{1,0}, f32[4096,27]{1,0}) tuple(f32[4096,27]{1,0} %pad.151, f32[4096,27]{1,0} %pad.101.clone.1)
}

%fused_computation.165 (param_0.756: f32[4096,55], param_1.996: f32[4096,27], param_2.1046: f32[4096,27], param_3.922: f32[4096,221]) -> (f32[4096,110], f32[4096,110]) {
  %param_0.756 = f32[4096,55]{1,0} parameter(0)
  %slice.694 = f32[4096,1]{1,0} slice(f32[4096,55]{1,0} %param_0.756), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_1.996 = f32[4096,27]{1,0} parameter(1)
  %param_2.1046 = f32[4096,27]{1,0} parameter(2)
  %add.991 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %param_1.996, f32[4096,27]{1,0} %param_2.1046), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.693 = f32[4096,27]{1,0} slice(f32[4096,55]{1,0} %param_0.756), slice={[0:4096:1], [2:55:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 55) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.990 = f32[4096,27]{1,0} add(f32[4096,27]{1,0} %add.991, f32[4096,27]{1,0} %slice.693), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.259 = f32[4096,28]{1,0} concatenate(f32[4096,1]{1,0} %slice.694, f32[4096,27]{1,0} %add.990), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_1770 = f32[] constant(0)
  %pad.164 = f32[4096,55]{1,0} pad(f32[4096,28]{1,0} %concatenate.259, f32[] %constant_1770), padding=0_0_0x0_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.163 = f32[4096,55]{1,0} pad(f32[4096,27]{1,0} %add.991, f32[] %constant_1770), padding=0_0_0x1_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.989 = f32[4096,55]{1,0} add(f32[4096,55]{1,0} %pad.164, f32[4096,55]{1,0} %pad.163), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.160 = f32[4096,110]{1,0} pad(f32[4096,55]{1,0} %add.989, f32[] %constant_1770), padding=0_0_0x1_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %param_3.922 = f32[4096,221]{1,0} parameter(3)
  %slice.658.clone.1 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_3.922), slice={[0:4096:1], [2:221:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 220) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.17.clone.1 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.658.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 220) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.657.clone.1 = f32[4096,110]{1,0} slice(f32[4096,221]{1,0} %param_3.922), slice={[0:4096:1], [1:220:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %reverse.16.clone.1 = f32[4096,110]{1,0} reverse(f32[4096,110]{1,0} %slice.657.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 1) limit_indices=(4096, 221) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.955.clone.1 = f32[4096,110]{1,0} add(f32[4096,110]{1,0} %reverse.17.clone.1, f32[4096,110]{1,0} %reverse.16.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.279.clone.1 = f32[4096,1]{1,0} slice(f32[4096,110]{1,0} %add.955.clone.1), slice={[0:4096], [0:1]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 1) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.278.clone.1 = f32[4096,54]{1,0} slice(f32[4096,55]{1,0} %add.989), slice={[0:4096], [0:54]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 0) limit_indices=(4096, 54) strides=(1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %slice.277.clone.1 = f32[4096,54]{1,0} slice(f32[4096,110]{1,0} %add.955.clone.1), slice={[0:4096:1], [2:110:2]}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/slice[start_indices=(0, 2) limit_indices=(4096, 110) strides=(1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %add.346.clone.1 = f32[4096,54]{1,0} add(f32[4096,54]{1,0} %slice.278.clone.1, f32[4096,54]{1,0} %slice.277.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %concatenate.102.clone.1 = f32[4096,55]{1,0} concatenate(f32[4096,1]{1,0} %slice.279.clone.1, f32[4096,54]{1,0} %add.346.clone.1), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.98.clone.1 = f32[4096,110]{1,0} pad(f32[4096,55]{1,0} %concatenate.102.clone.1, f32[] %constant_1770), padding=0_0_0x0_1_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  ROOT %tuple.49 = (f32[4096,110]{1,0}, f32[4096,110]{1,0}) tuple(f32[4096,110]{1,0} %pad.160, f32[4096,110]{1,0} %pad.98.clone.1)
}

%fused_computation.166 (param_0.755: s32[]) -> (f32[], f32[], s32[]) {
  %param_0.755 = s32[] parameter(0)
  %constant_691_clone_1 = s32[] constant(0)
  %compare.131.clone.1 = pred[] compare(s32[] %param_0.755, s32[] %constant_691_clone_1), direction=LE, metadata={op_name="jit(training_step)/jit(main)/le" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=199}
  %constant_693_clone_1 = f32[] constant(1)
  %constant_694_clone_1 = f32[] constant(0.1)
  %convert.19.clone.1 = f32[] convert(s32[] %param_0.755), metadata={op_name="jit(training_step)/jit(main)/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=195}
  %constant_455_clone_1 = f32[] constant(3.33333337e-05)
  %multiply.346.clone.1 = f32[] multiply(f32[] %convert.19.clone.1, f32[] %constant_455_clone_1), metadata={op_name="jit(training_step)/jit(main)/div" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=195}
  %power.2.clone.1 = f32[] power(f32[] %constant_694_clone_1, f32[] %multiply.346.clone.1), metadata={op_name="jit(training_step)/jit(main)/pow" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=199}
  %select.124.clone.1 = f32[] select(pred[] %compare.131.clone.1, f32[] %constant_693_clone_1, f32[] %power.2.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(_where)/jit(jit__where)/select_n" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=198}
  %maximum.60.clone.1 = f32[] maximum(f32[] %select.124.clone.1, f32[] %constant_694_clone_1), metadata={op_name="jit(training_step)/jit(main)/max" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=201}
  %constant_1813 = f32[] constant(0.001)
  %multiply.619 = f32[] multiply(f32[] %maximum.60.clone.1, f32[] %constant_1813), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=356}
  %constant_1812 = s32[] constant(1)
  %add.1016 = s32[] add(s32[] %param_0.755, s32[] %constant_1812), metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=364}
  ROOT %tuple.69 = (f32[], f32[], s32[]) tuple(f32[] %multiply.619, f32[] %maximum.60.clone.1, s32[] %add.1016)
}

%fused_computation.167 (param_0.753: f32[4096,3], param_1.1000: f32[4096,3]) -> (pred[4096], f32[4096], f32[4096]) {
  %param_1.1000 = f32[4096,3]{1,0} parameter(1)
  %constant_683_clone_1 = f32[] constant(-inf)
  %reduce.47.clone.1 = f32[4096]{0} reduce(f32[4096,3]{1,0} %param_1.1000, f32[] %constant_683_clone_1), dimensions={1}, to_apply=%region_2.229, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_max[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %constant_685_clone_1 = f32[] constant(0)
  %broadcast.1105.clone.1 = f32[4096]{0} broadcast(f32[] %constant_685_clone_1), dimensions={}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %maximum.65.clone.1 = f32[4096]{0} maximum(f32[4096]{0} %reduce.47.clone.1, f32[4096]{0} %broadcast.1105.clone.1), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/max" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=372}
  %param_0.753 = f32[4096,3]{1,0} parameter(0)
  %constant_1818 = f32[] constant(inf)
  %reduce.50 = f32[4096]{0} reduce(f32[4096,3]{1,0} %param_0.753, f32[] %constant_1818), dimensions={1}, to_apply=%region_3.233, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_min[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=373}
  %compare.352 = pred[4096]{0} compare(f32[4096]{0} %maximum.65.clone.1, f32[4096]{0} %reduce.50), direction=LT, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=375}
  ROOT %tuple.82 = (pred[4096]{0}, f32[4096]{0}, f32[4096]{0}) tuple(pred[4096]{0} %compare.352, f32[4096]{0} %reduce.50, f32[4096]{0} %maximum.65.clone.1)
}

%input_fused_computation_reduce.6 (param_0.763: f32[368,128], param_1.1005: f32[], param_2.1053: f32[368,128]) -> (f32[128], f32[128]) {
  %param_0.763 = f32[368,128]{1,0} parameter(0)
  %param_1.1005 = f32[] parameter(1)
  %reduce.51 = f32[128]{0} reduce(f32[368,128]{1,0} %param_0.763, f32[] %param_1.1005), dimensions={0}, to_apply=%region_28.1721, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  %param_2.1053 = f32[368,128]{1,0} parameter(2)
  %reduce.52.clone.1 = f32[128]{0} reduce(f32[368,128]{1,0} %param_2.1053, f32[] %param_1.1005), dimensions={0}, to_apply=%region_29.1725, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  ROOT %tuple.89 = (f32[128]{0}, f32[128]{0}) tuple(f32[128]{0} %reduce.51, f32[128]{0} %reduce.52.clone.1)
}

ENTRY %main.2269 (Arg_0.1: f32[144,27], Arg_1.2: f32[128], Arg_2.3: f32[150,128], Arg_3.4: f32[128], Arg_4.5: f32[128,128], Arg_5.6: f32[3], Arg_6.7: f32[128,3], Arg_7.8: f32[3,48,128], Arg_8.9: f32[3,48,128,128], Arg_9.10: f32[3,16,128], Arg_10.11: f32[3,16,128,128], Arg_11.12: s32[], Arg_12.13: f32[144,27], Arg_13.14: f32[128], Arg_14.15: f32[150,128], Arg_15.16: f32[128], Arg_16.17: f32[128,128], Arg_17.18: f32[3], Arg_18.19: f32[128,3], Arg_19.20: f32[3,48,128], Arg_20.21: f32[3,48,128,128], Arg_21.22: f32[3,16,128], Arg_22.23: f32[3,16,128,128], Arg_23.24: f32[144,27], Arg_24.25: f32[128], Arg_25.26: f32[150,128], Arg_26.27: f32[128], Arg_27.28: f32[128,128], Arg_28.29: f32[3], Arg_29.30: f32[128,3], Arg_30.31: f32[3,48,128], Arg_31.32: f32[3,48,128,128], Arg_32.33: f32[3,16,128], Arg_33.34: f32[3,16,128,128], Arg_34.35: f32[2,3], Arg_35.36: u32[2], Arg_36.37: s32[], Arg_37.38: f32[4096,3], Arg_38.39: f32[4096,3], Arg_39.40: f32[4096,3]) -> (f32[144,27], f32[128], f32[150,128], f32[128], f32[128,128], /*index=5*/f32[3], f32[128,3], f32[3,48,128], f32[3,48,128,128], f32[3,16,128], /*index=10*/f32[3,16,128,128], s32[], f32[144,27], f32[128], f32[150,128], /*index=15*/f32[128], f32[128,128], f32[3], f32[128,3], f32[3,48,128], /*index=20*/f32[3,48,128,128], f32[3,16,128], f32[3,16,128,128], f32[144,27], f32[128], /*index=25*/f32[150,128], f32[128], f32[128,128], f32[3], f32[128,3], /*index=30*/f32[3,48,128], f32[3,48,128,128], f32[3,16,128], f32[3,16,128,128], f32[2,3], /*index=35*/u32[2], s32[], f32[], f32[], f32[], /*index=40*/f32[]) {
  %Arg_0.1 = f32[144,27]{1,0} parameter(0)
  %copy.41 = f32[144,27]{1,0} copy(f32[144,27]{1,0} %Arg_0.1)
  %Arg_11.12 = s32[] parameter(11)
  %fusion.56 = (f32[], f32[], s32[]) fusion(s32[] %Arg_11.12), kind=kLoop, calls=%fused_computation.56, metadata={op_name="jit(training_step)/jit(main)/sub" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/transform.py" source_line=86}
  %get-tuple-element.186 = f32[] get-tuple-element((f32[], f32[], s32[]) %fusion.56), index=0
  %get-tuple-element.187 = f32[] get-tuple-element((f32[], f32[], s32[]) %fusion.56), index=1
  %Arg_36.37 = s32[] parameter(36), metadata={op_name="jit(training_step)/jit(main)/sub" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/schedule.py" source_line=194}
  %fusion.166 = (f32[], f32[], s32[]) fusion(s32[] %Arg_36.37), kind=kLoop, calls=%fused_computation.166, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=356}
  %get-tuple-element.184 = f32[] get-tuple-element((f32[], f32[], s32[]) %fusion.166), index=1
  %Arg_23.24 = f32[144,27]{1,0} parameter(23)
  %Arg_7.8 = f32[3,48,128]{2,1,0} parameter(7)
  %copy.48 = f32[3,48,128]{2,1,0} copy(f32[3,48,128]{2,1,0} %Arg_7.8)
  %constant_367 = s32[3,1]{1,0} constant({ {1}, {2}, {0} }), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(3, 1) broadcast_dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=59}
  %constant_363 = s32[3,1]{1,0} constant({ {2}, {0}, {1} }), metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(3, 1) broadcast_dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=58}
  %Arg_34.35 = f32[2,3]{1,0} parameter(34)
  %Arg_35.36 = u32[2]{0} parameter(35)
  %fusion.149 = (u32[2]{0}, u32[2]{0}) fusion(u32[2]{0} %Arg_35.36), kind=kLoop, calls=%fused_computation.149, metadata={op_name="jit(training_step)/jit(main)/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %get-tuple-element.216 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}) %fusion.149), index=0
  %get-tuple-element.217 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}) %fusion.149), index=1
  %fusion.156 = u32[2]{0} fusion(), kind=kLoop, calls=%fused_computation.156, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(0,) limit_indices=(2,) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %fusion.155 = u32[2]{0} fusion(), kind=kLoop, calls=%fused_computation.155, metadata={op_name="jit(training_step)/jit(main)/slice[start_indices=(2,) limit_indices=(4,) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %custom-call.212 = (u32[2]{0}, u32[2]{0}) custom-call(u32[2]{0} %get-tuple-element.216, u32[2]{0} %get-tuple-element.217, u32[2]{0} %fusion.156, u32[2]{0} %fusion.155), custom_call_target="cuda_threefry2x32", operand_layout_constraints={u32[2]{0}, u32[2]{0}, u32[2]{0}, u32[2]{0}}, api_version=API_VERSION_STATUS_RETURNING, metadata={op_name="jit(training_step)/jit(main)/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}, backend_config="\002\000\000\000\000\000\000\000"
  %get-tuple-element.213 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}) %custom-call.212), index=0, metadata={op_name="jit(training_step)/jit(main)/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %get-tuple-element.214 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}) %custom-call.212), index=1, metadata={op_name="jit(training_step)/jit(main)/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %fusion.2 = (u32[2]{0}, u32[2]{0}, u32[2]{0}) fusion(u32[2]{0} %get-tuple-element.213, u32[2]{0} %get-tuple-element.214), kind=kLoop, calls=%fused_computation.2, metadata={op_name="jit(training_step)/jit(main)/squeeze[dimensions=(0,)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=236}
  %get-tuple-element.215 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}, u32[2]{0}) %fusion.2), index=2
  %get-tuple-element.212 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}, u32[2]{0}) %fusion.2), index=1
  %custom-call.0 = (u32[2]{0}, u32[2]{0}) custom-call(u32[2]{0} %get-tuple-element.215, u32[2]{0} %get-tuple-element.212, u32[2]{0} %fusion.156, u32[2]{0} %fusion.155), custom_call_target="cuda_threefry2x32", operand_layout_constraints={u32[2]{0}, u32[2]{0}, u32[2]{0}, u32[2]{0}}, api_version=API_VERSION_STATUS_RETURNING, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}, backend_config="\002\000\000\000\000\000\000\000"
  %get-tuple-element.29 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}) %custom-call.0), index=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %get-tuple-element.30 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}) %custom-call.0), index=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=112}
  %fusion.139 = (u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}) fusion(u32[2]{0} %get-tuple-element.29, u32[2]{0} %get-tuple-element.30), kind=kLoop, calls=%fused_computation.139, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %get-tuple-element.210 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}) %fusion.139), index=3
  %get-tuple-element.209 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}) %fusion.139), index=2
  %fusion.154 = u32[111]{0} fusion(), kind=kLoop, calls=%fused_computation.154, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(0,) limit_indices=(111,) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %fusion.143 = u32[111]{0} fusion(), kind=kLoop, calls=%fused_computation.143, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/slice[start_indices=(111,) limit_indices=(222,) strides=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %custom-call.2 = (u32[111]{0}, u32[111]{0}) custom-call(u32[111]{0} %get-tuple-element.210, u32[111]{0} %get-tuple-element.209, u32[111]{0} %fusion.154, u32[111]{0} %fusion.143), custom_call_target="cuda_threefry2x32", operand_layout_constraints={u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}}, api_version=API_VERSION_STATUS_RETURNING, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}, backend_config="o\000\000\000\000\000\000\000"
  %get-tuple-element.33 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}) %custom-call.2), index=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %get-tuple-element.34 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}) %custom-call.2), index=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=325}
  %Arg_39.40 = f32[4096,3]{1,0} parameter(39)
  %Arg_38.39 = f32[4096,3]{1,0} parameter(38)
  %fusion.159 = (f32[4096,3]{1,0}, f32[4096,3]{1,0}) fusion(f32[4096,3]{1,0} %Arg_39.40, f32[4096,3]{1,0} %Arg_38.39, f32[2,3]{1,0} %Arg_34.35), kind=kLoop, calls=%fused_computation.159, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_max[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=368}
  %get-tuple-element.207 = f32[4096,3]{1,0} get-tuple-element((f32[4096,3]{1,0}, f32[4096,3]{1,0}) %fusion.159), index=0
  %get-tuple-element.208 = f32[4096,3]{1,0} get-tuple-element((f32[4096,3]{1,0}, f32[4096,3]{1,0}) %fusion.159), index=1
  %fusion.167 = (pred[4096]{0}, f32[4096]{0}, f32[4096]{0}) fusion(f32[4096,3]{1,0} %get-tuple-element.207, f32[4096,3]{1,0} %get-tuple-element.208), kind=kLoop, calls=%fused_computation.167, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/lt" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=375}
  %get-tuple-element.204 = pred[4096]{0} get-tuple-element((pred[4096]{0}, f32[4096]{0}, f32[4096]{0}) %fusion.167), index=0
  %get-tuple-element.206 = f32[4096]{0} get-tuple-element((pred[4096]{0}, f32[4096]{0}, f32[4096]{0}) %fusion.167), index=2
  %get-tuple-element.205 = f32[4096]{0} get-tuple-element((pred[4096]{0}, f32[4096]{0}, f32[4096]{0}) %fusion.167), index=1
  %fusion.141 = f32[3,4096,221]{2,1,0} fusion(f32[2,3]{1,0} %Arg_34.35, u32[111]{0} %get-tuple-element.33, u32[111]{0} %get-tuple-element.34, f32[4096,3]{1,0} %Arg_39.40, f32[4096,3]{1,0} %Arg_38.39, /*index=5*/pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, f32[4096]{0} %get-tuple-element.205), kind=kLoop, calls=%fused_computation.141, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=128}
  %get-tuple-element.192 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}) %fusion.139), index=0
  %get-tuple-element.193 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}) %fusion.139), index=1
  %custom-call.1 = (u32[111]{0}, u32[111]{0}) custom-call(u32[111]{0} %get-tuple-element.192, u32[111]{0} %get-tuple-element.193, u32[111]{0} %fusion.154, u32[111]{0} %fusion.143), custom_call_target="cuda_threefry2x32", operand_layout_constraints={u32[111]{0}, u32[111]{0}, u32[111]{0}, u32[111]{0}}, api_version=API_VERSION_STATUS_RETURNING, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}, backend_config="o\000\000\000\000\000\000\000"
  %get-tuple-element.31 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}) %custom-call.1), index=0, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %get-tuple-element.32 = u32[111]{0} get-tuple-element((u32[111]{0}, u32[111]{0}) %custom-call.1), index=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/threefry2x32" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %fusion.137 = f32[221]{0} fusion(u32[111]{0} %get-tuple-element.31, u32[111]{0} %get-tuple-element.32), kind=kLoop, calls=%fused_computation.137, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/log" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %Arg_9.10 = f32[3,16,128]{2,1,0} parameter(9)
  %copy.50 = f32[3,16,128]{2,1,0} copy(f32[3,16,128]{2,1,0} %Arg_9.10)
  %fusion.130 = (f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}) fusion(f32[3,4096,221]{2,1,0} %fusion.141, s32[3,1]{1,0} %constant_367, s32[3,1]{1,0} %constant_363), kind=kLoop, calls=%fused_computation.130, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/jit(jvp(vmap(vmap(vmap(_map_coordinates)))))/jit(jit_jvp(vmap(vmap(vmap(_map_coordinates)))))/squeeze[dimensions=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %get-tuple-element.202 = f32[3,4096,221]{2,1,0} get-tuple-element((f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}) %fusion.130), index=1
  %Arg_10.11 = f32[3,16,128,128]{3,2,1,0} parameter(10)
  %copy.51 = f32[3,16,128,128]{3,2,1,0} copy(f32[3,16,128,128]{3,2,1,0} %Arg_10.11)
  %get-tuple-element.203 = f32[3,4096,221]{2,1,0} get-tuple-element((f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}) %fusion.130), index=2
  %get-tuple-element.201 = f32[3,4096,221]{2,1,0} get-tuple-element((f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}, f32[3,4096,221]{2,1,0}) %fusion.130), index=0
  %fusion.115 = f32[4096,221]{1,0} fusion(f32[3,16,128]{2,1,0} %copy.50, f32[3,4096,221]{2,1,0} %get-tuple-element.202, f32[3,16,128,128]{3,2,1,0} %copy.51, f32[3,4096,221]{2,1,0} %get-tuple-element.203, f32[3,4096,221]{2,1,0} %get-tuple-element.201), kind=kLoop, calls=%fused_computation.115, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %fusion.114 = f32[4096,221]{1,0} fusion(f32[4096,221]{1,0} %fusion.115), kind=kLoop, calls=%fused_computation.114, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(softplus))/jit(jit_jvp(softplus))/jit(logaddexp)/jit(jit_logaddexp)/select_n" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=156}
  %fusion.110 = f32[4096,55]{1,0} fusion(pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %fusion.114), kind=kLoop, calls=%fused_computation.110, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.108 = f32[4096,13]{1,0} fusion(f32[4096,55]{1,0} %fusion.110), kind=kLoop, calls=%fused_computation.108, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.106 = f32[4096,3]{1,0} fusion(f32[4096,13]{1,0} %fusion.108), kind=kLoop, calls=%fused_computation.106, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.105 = f32[4096,6]{1,0} fusion(f32[4096,3]{1,0} %fusion.106, f32[4096,13]{1,0} %fusion.108), kind=kLoop, calls=%fused_computation.105, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (0, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %constant_70 = f32[] constant(0)
  %pad.59 = f32[4096,6]{1,0} pad(f32[4096,3]{1,0} %fusion.106, f32[] %constant_70), padding=0_0_0x1_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.161 = (f32[4096,27]{1,0}, f32[4096,27]{1,0}) fusion(f32[4096,13]{1,0} %fusion.108, f32[4096,6]{1,0} %fusion.105, f32[4096,6]{1,0} %pad.59, f32[4096,55]{1,0} %fusion.110), kind=kLoop, calls=%fused_computation.161, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %get-tuple-element.200 = f32[4096,27]{1,0} get-tuple-element((f32[4096,27]{1,0}, f32[4096,27]{1,0}) %fusion.161), index=1
  %get-tuple-element.199 = f32[4096,27]{1,0} get-tuple-element((f32[4096,27]{1,0}, f32[4096,27]{1,0}) %fusion.161), index=0
  %fusion.162 = (f32[4096,110]{1,0}, f32[4096,110]{1,0}) fusion(f32[4096,55]{1,0} %fusion.110, f32[4096,27]{1,0} %get-tuple-element.200, f32[4096,27]{1,0} %get-tuple-element.199, pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, /*index=5*/f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %fusion.114), kind=kLoop, calls=%fused_computation.162, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(_cumulative_reduction))/jit(jit_jvp(_cumulative_reduction))/pad[padding_config=((0, 0, 0), (1, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %get-tuple-element.198 = f32[4096,110]{1,0} get-tuple-element((f32[4096,110]{1,0}, f32[4096,110]{1,0}) %fusion.162), index=1
  %get-tuple-element.197 = f32[4096,110]{1,0} get-tuple-element((f32[4096,110]{1,0}, f32[4096,110]{1,0}) %fusion.162), index=0
  %fusion.98 = (f32[4096,221]{1,0}, f32[4096,221]{1,0}, f32[4096,221]{1,0}) fusion(f32[221]{0} %fusion.137, pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %fusion.114, /*index=5*/f32[4096,110]{1,0} %get-tuple-element.198, f32[4096,110]{1,0} %get-tuple-element.197), kind=kLoop, calls=%fused_computation.98, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/sub" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %get-tuple-element.194 = f32[4096,221]{1,0} get-tuple-element((f32[4096,221]{1,0}, f32[4096,221]{1,0}, f32[4096,221]{1,0}) %fusion.98), index=0
  %iota.62 = s32[4096,221]{1,0} iota(), iota_dimension=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/broadcast_in_dim[shape=(4096, 221) broadcast_dimensions=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %sort.2 = (f32[4096,221]{1,0}, s32[4096,221]{1,0}, s32[4096,221]{1,0}) sort(f32[4096,221]{1,0} %get-tuple-element.194, s32[4096,221]{1,0} %iota.62, s32[4096,221]{1,0} %iota.62), dimensions={1}, is_stable=true, to_apply=%region_5.628.clone, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/sort[dimension=1 is_stable=True num_keys=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %get-tuple-element.51 = s32[4096,221]{1,0} get-tuple-element((f32[4096,221]{1,0}, s32[4096,221]{1,0}, s32[4096,221]{1,0}) %sort.2), index=1, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(vmap(argsort)))/jit(jit_jvp(vmap(argsort)))/sort[dimension=1 is_stable=True num_keys=1]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=410}
  %fusion.97 = s32[4096,33,2]{1,0,2} fusion(s32[4096,221]{1,0} %get-tuple-element.51), kind=kLoop, calls=%fused_computation.97, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=2]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=420}
  %fusion.96 = f32[3,4096,3,33]{3,1,0,2} fusion(s32[3,1]{1,0} %constant_367, s32[3,1]{1,0} %constant_363, f32[3,4096,221]{2,1,0} %fusion.141, s32[4096,33,2]{1,0,2} %fusion.97), kind=kLoop, calls=%fused_computation.96, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/transpose[permutation=(0, 2, 1, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=165}
  %Arg_8.9 = f32[3,48,128,128]{3,2,1,0} parameter(8)
  %copy.49 = f32[3,48,128,128]{3,2,1,0} copy(f32[3,48,128,128]{3,2,1,0} %Arg_8.9)
  %fusion.75 = f32[135168,144]{0,1} fusion(f32[3,48,128]{2,1,0} %copy.48, f32[3,4096,3,33]{3,1,0,2} %fusion.96, f32[3,48,128,128]{3,2,1,0} %copy.49), kind=kLoop, calls=%fused_computation.75, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reshape[new_sizes=(135168, 144) dimensions=None]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=432}
  %cublas-gemm.1 = f32[135168,27]{1,0} custom-call(f32[135168,144]{0,1} %fusion.75, f32[144,27]{1,0} %copy.41), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=None]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"19464192\",\"rhs_stride\":\"3888\",\"selected_algorithm\":\"2\"}"
  %fusion.74 = f32[2]{0} fusion(), kind=kLoop, calls=%fused_computation.74, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/pow" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=26}
  %fusion.73 = f32[135168,27,4]{2,1,0} fusion(f32[135168,27]{1,0} %cublas-gemm.1, f32[2]{0} %fusion.74), kind=kLoop, calls=%fused_computation.73, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=2]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=31}
  %input_fusion_scatter.14 = u32[4096,221,3]{1,0,2} fusion(s32[4096,33,2]{1,0,2} %fusion.97), kind=kInput, calls=%input_fused_computation_scatter.14, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/scatter[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 1), scatter_dims_to_operand_dims=(0, 1)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %fusion.70 = pred[4096,33,3]{1,0,2} fusion(u32[4096,221,3]{1,0,2} %input_fusion_scatter.14, s32[4096,33,2]{1,0,2} %fusion.97), kind=kLoop, calls=%fused_computation.70, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/eq" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %get-tuple-element.195 = f32[4096,221]{1,0} get-tuple-element((f32[4096,221]{1,0}, f32[4096,221]{1,0}, f32[4096,221]{1,0}) %fusion.98), index=1
  %fusion.72 = f32[135168,150]{1,0} fusion(f32[135168,27]{1,0} %cublas-gemm.1, f32[2]{0} %fusion.74, f32[135168,27,4]{2,1,0} %fusion.73, f32[4096,3]{1,0} %Arg_39.40), kind=kLoop, calls=%fused_computation.72, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/concatenate[dimension=1]" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=73}
  %Arg_2.3 = f32[150,128]{1,0} parameter(2)
  %copy.43 = f32[150,128]{1,0} copy(f32[150,128]{1,0} %Arg_2.3)
  %Arg_1.2 = f32[128]{0} parameter(1)
  %broadcast.545 = f32[135168,128]{1,0} broadcast(f32[128]{0} %Arg_1.2), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  %cublas-gemm.5 = f32[135168,128]{1,0} custom-call(f32[135168,150]{1,0} %fusion.72, f32[150,128]{1,0} %copy.43, f32[135168,128]{1,0} %broadcast.545), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":1,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"20275200\",\"rhs_stride\":\"19200\",\"selected_algorithm\":\"105\"}"
  %fusion.153 = f32[135168,128]{1,0} fusion(f32[135168,128]{1,0} %cublas-gemm.5), kind=kLoop, calls=%fused_computation.153, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(relu)/jit(jit_relu)/max" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  %Arg_4.5 = f32[128,128]{1,0} parameter(4)
  %copy.45 = f32[128,128]{1,0} copy(f32[128,128]{1,0} %Arg_4.5)
  %Arg_3.4 = f32[128]{0} parameter(3)
  %broadcast.548 = f32[135168,128]{1,0} broadcast(f32[128]{0} %Arg_3.4), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  %cublas-gemm.9 = f32[135168,128]{1,0} custom-call(f32[135168,128]{1,0} %fusion.153, f32[128,128]{1,0} %copy.45, f32[135168,128]{1,0} %broadcast.548), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":1,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"17301504\",\"rhs_stride\":\"16384\",\"selected_algorithm\":\"4\"}"
  %fusion.152 = f32[135168,128]{1,0} fusion(f32[135168,128]{1,0} %cublas-gemm.9), kind=kLoop, calls=%fused_computation.152, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/jit(relu)/jit(jit_relu)/max" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=105}
  %Arg_6.7 = f32[128,3]{1,0} parameter(6)
  %copy.47 = f32[128,3]{1,0} copy(f32[128,3]{1,0} %Arg_6.7)
  %Arg_5.6 = f32[3]{0} parameter(5)
  %broadcast.551 = f32[135168,3]{0,1} broadcast(f32[3]{0} %Arg_5.6), dimensions={1}, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  %cublas-gemm.13 = f32[135168,3]{0,1} custom-call(f32[135168,128]{1,0} %fusion.152, f32[128,3]{1,0} %copy.47, f32[135168,3]{0,1} %broadcast.551), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":1,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"17301504\",\"rhs_stride\":\"384\",\"selected_algorithm\":\"3\"}"
  %Arg_37.38 = f32[4096,3]{1,0} parameter(37)
  %input_fusion_scatter.13 = f32[4096,221,3]{1,0,2} fusion(s32[4096,33,2]{1,0,2} %fusion.97, pred[4096,33,3]{1,0,2} %fusion.70, f32[135168,3]{0,1} %cublas-gemm.13), kind=kInput, calls=%input_fused_computation_scatter.13, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 1), scatter_dims_to_operand_dims=(0, 1)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=452}
  %input_fusion_reduce.5 = f32[3,4096]{0,1} fusion(f32[4096,221,3]{1,0,2} %input_fusion_scatter.13, f32[4096,221]{1,0} %get-tuple-element.195), kind=kInput, calls=%input_fused_computation_reduce.5
  %get-tuple-element.196 = f32[4096,221]{1,0} get-tuple-element((f32[4096,221]{1,0}, f32[4096,221]{1,0}, f32[4096,221]{1,0}) %fusion.98), index=2
  %input_fusion_reduce.4 = f32[4096]{0} fusion(f32[4096,221]{1,0} %get-tuple-element.195, s32[4096,33,2]{1,0,2} %fusion.97), kind=kInput, calls=%input_fused_computation_reduce.4, metadata={op_name="jit(training_step)/jit(main)/jit(jvp(render_rays))/jit(jit_jvp(render_rays))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=486}
  %fusion.66 = f32[4096,3]{1,0} fusion(f32[4096,3]{1,0} %Arg_37.38, f32[3,4096]{0,1} %input_fusion_reduce.5, f32[4096,221]{1,0} %get-tuple-element.196, f32[4096]{0} %input_fusion_reduce.4), kind=kLoop, calls=%fused_computation.66, metadata={op_name="jit(training_step)/jit(main)/sub" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %fusion.64 = f32[135168,3]{0,1} fusion(pred[4096,33,3]{1,0,2} %fusion.70, s32[4096,33,2]{1,0,2} %fusion.97, f32[4096,221]{1,0} %get-tuple-element.195, f32[135168,3]{0,1} %cublas-gemm.13, f32[4096,3]{1,0} %fusion.66, /*index=5*/f32[4096]{0} %input_fusion_reduce.4, f32[4096,221]{1,0} %get-tuple-element.196), kind=kLoop, calls=%fused_computation.64, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(sigmoid)))/jit(jit_transpose(jvp(sigmoid)))/mul" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=116}
  %cublas-gemm.15 = f32[135168,128]{1,0} custom-call(f32[135168,3]{0,1} %fusion.64, f32[128,3]{1,0} %copy.47), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/dot_general[dimension_numbers=(((1,), (1,)), ((), ())) precision=None preferred_element_type=None]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"1\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"405504\",\"rhs_stride\":\"384\",\"selected_algorithm\":\"0\"}"
  %fusion.63 = f32[135168,128]{1,0} fusion(f32[135168,128]{1,0} %cublas-gemm.15, f32[135168,128]{1,0} %cublas-gemm.9), kind=kLoop, calls=%fused_computation.63, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=105}
  %cublas-gemm.17 = f32[135168,128]{1,0} custom-call(f32[135168,128]{1,0} %fusion.63, f32[128,128]{1,0} %copy.45), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/dot_general[dimension_numbers=(((1,), (1,)), ((), ())) precision=None preferred_element_type=None]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"1\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"17301504\",\"rhs_stride\":\"16384\",\"selected_algorithm\":\"14\"}"
  %fusion.62 = f32[135168,128]{1,0} fusion(f32[135168,128]{1,0} %cublas-gemm.17, f32[135168,128]{1,0} %cublas-gemm.5), kind=kLoop, calls=%fused_computation.62, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/select_n" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=96}
  %cublas-gemm.19 = f32[135168,150]{1,0} custom-call(f32[135168,128]{1,0} %fusion.62, f32[150,128]{1,0} %copy.43), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/dot_general[dimension_numbers=(((1,), (1,)), ((), ())) precision=None preferred_element_type=None]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"1\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"17301504\",\"rhs_stride\":\"19200\",\"selected_algorithm\":\"21\"}"
  %fusion.61 = f32[135168,27]{1,0} fusion(f32[135168,27,4]{2,1,0} %fusion.73, f32[135168,150]{1,0} %cublas-gemm.19, f32[2]{0} %fusion.74), kind=kLoop, calls=%fused_computation.61, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/add_any" source_file="/home/brent/tensorf-jax/tensorf/networks.py" source_line=27}
  %cublas-gemm.21 = f32[144,27]{1,0} custom-call(f32[135168,144]{0,1} %fusion.75, f32[135168,27]{1,0} %fusion.61), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/transpose[permutation=(1, 0)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"0\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"19464192\",\"rhs_stride\":\"3649536\",\"selected_algorithm\":\"0\"}"
  %Arg_12.13 = f32[144,27]{1,0} parameter(12)
  %fusion.54 = (f32[144,27]{1,0}, f32[144,27]{1,0}, f32[144,27]{1,0}) fusion(f32[144,27]{1,0} %copy.41, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[144,27]{1,0} %Arg_23.24, /*index=5*/f32[144,27]{1,0} %cublas-gemm.21, f32[144,27]{1,0} %Arg_12.13), kind=kLoop, calls=%fused_computation.54, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.222 = f32[144,27]{1,0} get-tuple-element((f32[144,27]{1,0}, f32[144,27]{1,0}, f32[144,27]{1,0}) %fusion.54), index=0
  %Arg_26.27 = f32[128]{0} parameter(26)
  %input_fusion_reduce.2 = (f32[368,128]{1,0}, f32[368,128]{1,0}) fusion(f32[135168,128]{1,0} %fusion.63, f32[135168,128]{1,0} %fusion.62), kind=kInput, calls=%input_fused_computation_reduce.2
  %get-tuple-element.220 = f32[368,128]{1,0} get-tuple-element((f32[368,128]{1,0}, f32[368,128]{1,0}) %input_fusion_reduce.2), index=0
  %get-tuple-element.221 = f32[368,128]{1,0} get-tuple-element((f32[368,128]{1,0}, f32[368,128]{1,0}) %input_fusion_reduce.2), index=1
  %input_fusion_reduce.6 = (f32[128]{0}, f32[128]{0}) fusion(f32[368,128]{1,0} %get-tuple-element.220, f32[] %constant_70, f32[368,128]{1,0} %get-tuple-element.221), kind=kInput, calls=%input_fused_computation_reduce.6, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  %get-tuple-element.218 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}) %input_fusion_reduce.6), index=0
  %Arg_15.16 = f32[128]{0} parameter(15)
  %Arg_24.25 = f32[128]{0} parameter(24)
  %get-tuple-element.219 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}) %input_fusion_reduce.6), index=1
  %Arg_13.14 = f32[128]{0} parameter(13)
  %fusion.45 = (f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) fusion(f32[128]{0} %Arg_3.4, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[128]{0} %Arg_26.27, /*index=5*/f32[128]{0} %get-tuple-element.218, f32[128]{0} %Arg_15.16, f32[128]{0} %Arg_1.2, f32[128]{0} %Arg_24.25, f32[128]{0} %get-tuple-element.219, /*index=10*/f32[128]{0} %Arg_13.14), kind=kLoop, calls=%fused_computation.45, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.223 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) %fusion.45), index=3
  %Arg_25.26 = f32[150,128]{1,0} parameter(25)
  %cublas-gemm.23 = f32[150,128]{1,0} custom-call(f32[135168,150]{1,0} %fusion.72, f32[135168,128]{1,0} %fusion.62), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/transpose[permutation=(1, 0)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"0\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"20275200\",\"rhs_stride\":\"17301504\",\"selected_algorithm\":\"-1\"}"
  %Arg_14.15 = f32[150,128]{1,0} parameter(14)
  %fusion.48 = (f32[150,128]{1,0}, f32[150,128]{1,0}, f32[150,128]{1,0}) fusion(f32[150,128]{1,0} %copy.43, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[150,128]{1,0} %Arg_25.26, /*index=5*/f32[150,128]{1,0} %cublas-gemm.23, f32[150,128]{1,0} %Arg_14.15), kind=kLoop, calls=%fused_computation.48, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.224 = f32[150,128]{1,0} get-tuple-element((f32[150,128]{1,0}, f32[150,128]{1,0}, f32[150,128]{1,0}) %fusion.48), index=0
  %get-tuple-element.225 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) %fusion.45), index=0
  %Arg_27.28 = f32[128,128]{1,0} parameter(27)
  %cublas-gemm.25 = f32[128,128]{1,0} custom-call(f32[135168,128]{1,0} %fusion.153, f32[135168,128]{1,0} %fusion.63), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/transpose[permutation=(1, 0)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"0\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"17301504\",\"rhs_stride\":\"17301504\",\"selected_algorithm\":\"4\"}"
  %Arg_16.17 = f32[128,128]{1,0} parameter(16)
  %fusion.42 = (f32[128,128]{1,0}, f32[128,128]{1,0}, f32[128,128]{1,0}) fusion(f32[128,128]{1,0} %copy.45, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[128,128]{1,0} %Arg_27.28, /*index=5*/f32[128,128]{1,0} %cublas-gemm.25, f32[128,128]{1,0} %Arg_16.17), kind=kLoop, calls=%fused_computation.42, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.226 = f32[128,128]{1,0} get-tuple-element((f32[128,128]{1,0}, f32[128,128]{1,0}, f32[128,128]{1,0}) %fusion.42), index=0
  %Arg_28.29 = f32[3]{0} parameter(28)
  %input_fusion_reduce.1 = f32[3,368]{1,0} fusion(f32[135168,3]{0,1} %fusion.64), kind=kInput, calls=%input_fused_computation_reduce.1
  %reduce.17 = f32[3]{0} reduce(f32[3,368]{1,0} %input_fusion_reduce.1, f32[] %constant_70), dimensions={1}, to_apply=%region_27.1717, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(0,)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=195}
  %Arg_17.18 = f32[3]{0} parameter(17)
  %fusion.39 = (f32[3]{0}, f32[3]{0}, f32[3]{0}) fusion(f32[3]{0} %Arg_5.6, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[3]{0} %Arg_28.29, /*index=5*/f32[3]{0} %reduce.17, f32[3]{0} %Arg_17.18), kind=kLoop, calls=%fused_computation.39, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.227 = f32[3]{0} get-tuple-element((f32[3]{0}, f32[3]{0}, f32[3]{0}) %fusion.39), index=0
  %Arg_29.30 = f32[128,3]{1,0} parameter(29)
  %cublas-gemm.27 = f32[128,3]{1,0} custom-call(f32[135168,128]{1,0} %fusion.152, f32[135168,3]{0,1} %fusion.64), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/transpose[permutation=(1, 0)]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"0\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"17301504\",\"rhs_stride\":\"405504\",\"selected_algorithm\":\"0\"}"
  %Arg_18.19 = f32[128,3]{1,0} parameter(18)
  %fusion.36 = (f32[128,3]{1,0}, f32[128,3]{1,0}, f32[128,3]{1,0}) fusion(f32[128,3]{1,0} %copy.47, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[128,3]{1,0} %Arg_29.30, /*index=5*/f32[128,3]{1,0} %cublas-gemm.27, f32[128,3]{1,0} %Arg_18.19), kind=kLoop, calls=%fused_computation.36, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.228 = f32[128,3]{1,0} get-tuple-element((f32[128,3]{1,0}, f32[128,3]{1,0}, f32[128,3]{1,0}) %fusion.36), index=0
  %Arg_30.31 = f32[3,48,128]{2,1,0} parameter(30)
  %cublas-gemm.29 = f32[135168,144]{1,0} custom-call(f32[135168,27]{1,0} %fusion.61, f32[144,27]{1,0} %copy.41), custom_call_target="__cublas$gemm", metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/dot_general[dimension_numbers=(((1,), (1,)), ((), ())) precision=None preferred_element_type=None]" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/flax/linen/linear.py" source_line=188}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"1\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"batch_size\":\"1\",\"lhs_stride\":\"3649536\",\"rhs_stride\":\"3888\",\"selected_algorithm\":\"0\"}"
  %fusion.34 = f32[3,4096,48,33]{2,1,0,3} fusion(f32[3,48,128,128]{3,2,1,0} %copy.49, f32[3,4096,3,33]{3,1,0,2} %fusion.96, f32[135168,144]{1,0} %cublas-gemm.29), kind=kLoop, calls=%fused_computation.34, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %input_fusion_scatter.12 = f32[3,48,128]{2,1,0} fusion(f32[3,4096,48,33]{2,1,0,3} %fusion.34, f32[3,4096,3,33]{3,1,0,2} %fusion.96), kind=kInput, calls=%input_fused_computation_scatter.12, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.11 = f32[3,48,128]{2,1,0} fusion(f32[3,4096,48,33]{2,1,0,3} %fusion.34, f32[3,4096,3,33]{3,1,0,2} %fusion.96), kind=kInput, calls=%input_fused_computation_scatter.11, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %Arg_19.20 = f32[3,48,128]{2,1,0} parameter(19)
  %fusion.31 = (f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}) fusion(f32[3,48,128]{2,1,0} %copy.48, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[3,48,128]{2,1,0} %Arg_30.31, /*index=5*/f32[3,48,128]{2,1,0} %input_fusion_scatter.12, f32[3,48,128]{2,1,0} %input_fusion_scatter.11, f32[3,48,128]{2,1,0} %Arg_19.20), kind=kLoop, calls=%fused_computation.31, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.229 = f32[3,48,128]{2,1,0} get-tuple-element((f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}) %fusion.31), index=0
  %Arg_31.32 = f32[3,48,128,128]{3,2,1,0} parameter(31)
  %fusion.30 = f32[3,4096,48,33]{3,2,1,0} fusion(f32[3,48,128]{2,1,0} %copy.48, f32[3,4096,3,33]{3,1,0,2} %fusion.96, f32[135168,144]{1,0} %cublas-gemm.29), kind=kLoop, calls=%fused_computation.30, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %input_fusion_scatter.10 = f32[3,48,128,128]{3,2,1,0} fusion(f32[3,4096,48,33]{3,2,1,0} %fusion.30, f32[3,4096,3,33]{3,1,0,2} %fusion.96), kind=kInput, calls=%input_fused_computation_scatter.10, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.9 = f32[3,48,128,128]{3,2,1,0} fusion(f32[3,4096,48,33]{3,2,1,0} %fusion.30, f32[3,4096,3,33]{3,1,0,2} %fusion.96), kind=kInput, calls=%input_fused_computation_scatter.9, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.8 = f32[3,48,128,128]{3,2,1,0} fusion(f32[3,4096,48,33]{3,2,1,0} %fusion.30, f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.10, f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.9, f32[3,4096,3,33]{3,1,0,2} %fusion.96), kind=kInput, calls=%input_fused_computation_scatter.8, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.7 = f32[3,48,128,128]{3,2,1,0} fusion(f32[3,4096,48,33]{3,2,1,0} %fusion.30, f32[3,4096,3,33]{3,1,0,2} %fusion.96), kind=kInput, calls=%input_fused_computation_scatter.7, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %Arg_20.21 = f32[3,48,128,128]{3,2,1,0} parameter(20)
  %fusion.27 = (f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}) fusion(f32[3,48,128,128]{3,2,1,0} %copy.49, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[3,48,128,128]{3,2,1,0} %Arg_31.32, /*index=5*/f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.8, f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.7, f32[3,48,128,128]{3,2,1,0} %Arg_20.21), kind=kLoop, calls=%fused_computation.27, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.230 = f32[3,48,128,128]{3,2,1,0} get-tuple-element((f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}) %fusion.27), index=0
  %Arg_32.33 = f32[3,16,128]{2,1,0} parameter(32)
  %fusion.24 = f32[4096,221]{1,0} fusion(f32[4096,221,3]{1,0,2} %input_fusion_scatter.13, f32[4096,3]{1,0} %fusion.66, f32[4096]{0} %input_fusion_reduce.4, f32[4096,221]{1,0} %get-tuple-element.196), kind=kLoop, calls=%fused_computation.24, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(2,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %fusion.25 = f32[4096]{0} fusion(f32[3,4096]{0,1} %input_fusion_reduce.5, f32[4096,3]{1,0} %fusion.66), kind=kLoop, calls=%fused_computation.25, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/reduce_sum[axes=(1,)]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=188}
  %input_fusion_scatter.6 = f32[4096,221]{1,0} fusion(f32[4096,221]{1,0} %fusion.24, s32[4096,33,2]{1,0,2} %fusion.97, f32[4096]{0} %fusion.25, f32[4096,221]{1,0} %get-tuple-element.196, f32[4096]{0} %input_fusion_reduce.4), kind=kInput, calls=%input_fused_computation_scatter.6, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(), inserted_window_dims=(0, 1), scatter_dims_to_operand_dims=(0, 1)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=472}
  %fusion.23 = f32[4096,221]{1,0} fusion(f32[4096,221]{1,0} %get-tuple-element.196, f32[4096,221]{1,0} %input_fusion_scatter.6, f32[4096]{0} %fusion.25, pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, /*index=5*/f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %fusion.114, f32[4096]{0} %input_fusion_reduce.4, f32[4096,3]{1,0} %fusion.66), kind=kLoop, calls=%fused_computation.23, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/mul" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.21 = f32[4096,55]{1,0} fusion(f32[4096,221]{1,0} %fusion.23), kind=kLoop, calls=%fused_computation.21, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.19 = f32[4096,13]{1,0} fusion(f32[4096,55]{1,0} %fusion.21), kind=kLoop, calls=%fused_computation.19, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.17 = f32[4096,3]{1,0} fusion(f32[4096,13]{1,0} %fusion.19), kind=kLoop, calls=%fused_computation.17, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/add" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.16 = f32[4096,6]{1,0} fusion(f32[4096,3]{1,0} %fusion.17, f32[4096,13]{1,0} %fusion.19), kind=kLoop, calls=%fused_computation.16, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (0, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %pad.73 = f32[4096,6]{1,0} pad(f32[4096,3]{1,0} %fusion.17, f32[] %constant_70), padding=0_0_0x1_0_1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %fusion.164 = (f32[4096,27]{1,0}, f32[4096,27]{1,0}) fusion(f32[4096,13]{1,0} %fusion.19, f32[4096,6]{1,0} %fusion.16, f32[4096,6]{1,0} %pad.73, f32[4096,55]{1,0} %fusion.21), kind=kLoop, calls=%fused_computation.164, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 1, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %get-tuple-element.155 = f32[4096,27]{1,0} get-tuple-element((f32[4096,27]{1,0}, f32[4096,27]{1,0}) %fusion.164), index=1
  %get-tuple-element.154 = f32[4096,27]{1,0} get-tuple-element((f32[4096,27]{1,0}, f32[4096,27]{1,0}) %fusion.164), index=0
  %fusion.165 = (f32[4096,110]{1,0}, f32[4096,110]{1,0}) fusion(f32[4096,55]{1,0} %fusion.21, f32[4096,27]{1,0} %get-tuple-element.155, f32[4096,27]{1,0} %get-tuple-element.154, f32[4096,221]{1,0} %fusion.23), kind=kLoop, calls=%fused_computation.165, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(_cumulative_reduction)))/jit(jit_transpose(jvp(_cumulative_reduction)))/pad[padding_config=((0, 0, 0), (1, 0, 1))]" source_file="/home/brent/tensorf-jax/tensorf/render.py" source_line=272}
  %get-tuple-element.153 = f32[4096,110]{1,0} get-tuple-element((f32[4096,110]{1,0}, f32[4096,110]{1,0}) %fusion.165), index=1
  %get-tuple-element.152 = f32[4096,110]{1,0} get-tuple-element((f32[4096,110]{1,0}, f32[4096,110]{1,0}) %fusion.165), index=0
  %fusion.10 = f32[3,4096,16,221]{2,1,0,3} fusion(f32[3,16,128,128]{3,2,1,0} %copy.51, f32[3,4096,221]{2,1,0} %get-tuple-element.203, f32[3,4096,221]{2,1,0} %get-tuple-element.201, f32[4096,221]{1,0} %fusion.114, f32[4096,221]{1,0} %fusion.115, /*index=5*/pred[4096]{0} %get-tuple-element.204, f32[4096]{0} %get-tuple-element.206, f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %input_fusion_scatter.6, f32[4096,221]{1,0} %get-tuple-element.196, /*index=10*/f32[4096,221]{1,0} %fusion.23, f32[4096,110]{1,0} %get-tuple-element.153, f32[4096,110]{1,0} %get-tuple-element.152), kind=kLoop, calls=%fused_computation.10, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %input_fusion_scatter.5 = f32[3,16,128]{2,1,0} fusion(f32[3,4096,16,221]{2,1,0,3} %fusion.10, f32[3,4096,221]{2,1,0} %get-tuple-element.202), kind=kInput, calls=%input_fused_computation_scatter.5, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.4 = f32[3,16,128]{2,1,0} fusion(f32[3,4096,16,221]{2,1,0,3} %fusion.10, f32[3,4096,221]{2,1,0} %get-tuple-element.202), kind=kInput, calls=%input_fused_computation_scatter.4, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2), scatter_dims_to_operand_dims=(0, 2)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %Arg_21.22 = f32[3,16,128]{2,1,0} parameter(21)
  %fusion.7 = (f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}) fusion(f32[3,16,128]{2,1,0} %copy.50, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[3,16,128]{2,1,0} %Arg_32.33, /*index=5*/f32[3,16,128]{2,1,0} %input_fusion_scatter.5, f32[3,16,128]{2,1,0} %input_fusion_scatter.4, f32[3,16,128]{2,1,0} %Arg_21.22), kind=kLoop, calls=%fused_computation.7, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.231 = f32[3,16,128]{2,1,0} get-tuple-element((f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}) %fusion.7), index=0
  %Arg_33.34 = f32[3,16,128,128]{3,2,1,0} parameter(33)
  %fusion.6 = f32[3,4096,16,221]{3,2,1,0} fusion(f32[3,16,128]{2,1,0} %copy.50, f32[3,4096,221]{2,1,0} %get-tuple-element.202, f32[4096,221]{1,0} %fusion.114, f32[4096,221]{1,0} %fusion.115, pred[4096]{0} %get-tuple-element.204, /*index=5*/f32[4096]{0} %get-tuple-element.206, f32[4096]{0} %get-tuple-element.205, f32[4096,221]{1,0} %input_fusion_scatter.6, f32[4096,221]{1,0} %get-tuple-element.196, f32[4096,221]{1,0} %fusion.23, /*index=10*/f32[4096,110]{1,0} %get-tuple-element.153, f32[4096,110]{1,0} %get-tuple-element.152), kind=kLoop, calls=%fused_computation.6, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %input_fusion_scatter.3 = f32[3,16,128,128]{3,2,1,0} fusion(f32[3,4096,16,221]{3,2,1,0} %fusion.6, f32[3,4096,221]{2,1,0} %get-tuple-element.203, f32[3,4096,221]{2,1,0} %get-tuple-element.201), kind=kInput, calls=%input_fused_computation_scatter.3, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.2 = f32[3,16,128,128]{3,2,1,0} fusion(f32[3,4096,16,221]{3,2,1,0} %fusion.6, f32[3,4096,221]{2,1,0} %get-tuple-element.203, f32[3,4096,221]{2,1,0} %get-tuple-element.201), kind=kInput, calls=%input_fused_computation_scatter.2, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter.1 = f32[3,16,128,128]{3,2,1,0} fusion(f32[3,4096,16,221]{3,2,1,0} %fusion.6, f32[3,16,128,128]{3,2,1,0} %input_fusion_scatter.3, f32[3,16,128,128]{3,2,1,0} %input_fusion_scatter.2, f32[3,4096,221]{2,1,0} %get-tuple-element.203, f32[3,4096,221]{2,1,0} %get-tuple-element.201), kind=kInput, calls=%input_fused_computation_scatter.1, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %input_fusion_scatter = f32[3,16,128,128]{3,2,1,0} fusion(f32[3,4096,16,221]{3,2,1,0} %fusion.6, f32[3,4096,221]{2,1,0} %get-tuple-element.203, f32[3,4096,221]{2,1,0} %get-tuple-element.201), kind=kInput, calls=%input_fused_computation_scatter, metadata={op_name="jit(training_step)/jit(main)/jit(transpose(jvp(render_rays)))/jit(jit_transpose(jvp(render_rays)))/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/jit(transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/jit(jit_transpose(jvp(vmap(vmap(vmap(_map_coordinates))))))/scatter-add[update_consts=() dimension_numbers=ScatterDimensionNumbers(update_window_dims=(2,), inserted_window_dims=(0, 2, 3), scatter_dims_to_operand_dims=(0, 2, 3)) indices_are_sorted=False unique_indices=False mode=GatherScatterMode.PROMISE_IN_BOUNDS]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %Arg_22.23 = f32[3,16,128,128]{3,2,1,0} parameter(22)
  %fusion.3 = (f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}) fusion(f32[3,16,128,128]{3,2,1,0} %copy.51, f32[] %get-tuple-element.186, f32[] %get-tuple-element.187, f32[] %get-tuple-element.184, f32[3,16,128,128]{3,2,1,0} %Arg_33.34, /*index=5*/f32[3,16,128,128]{3,2,1,0} %input_fusion_scatter.1, f32[3,16,128,128]{3,2,1,0} %input_fusion_scatter, f32[3,16,128,128]{3,2,1,0} %Arg_22.23), kind=kLoop, calls=%fused_computation.3, metadata={op_name="jit(training_step)/jit(main)/add" source_file="/home/brent/miniconda/envs/py38/lib/python3.8/site-packages/optax/_src/update.py" source_line=43}
  %get-tuple-element.232 = f32[3,16,128,128]{3,2,1,0} get-tuple-element((f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}) %fusion.3), index=0
  %get-tuple-element.233 = s32[] get-tuple-element((f32[], f32[], s32[]) %fusion.56), index=2
  %get-tuple-element.234 = f32[144,27]{1,0} get-tuple-element((f32[144,27]{1,0}, f32[144,27]{1,0}, f32[144,27]{1,0}) %fusion.54), index=2
  %get-tuple-element.235 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) %fusion.45), index=5
  %get-tuple-element.236 = f32[150,128]{1,0} get-tuple-element((f32[150,128]{1,0}, f32[150,128]{1,0}, f32[150,128]{1,0}) %fusion.48), index=2
  %get-tuple-element.237 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) %fusion.45), index=2
  %get-tuple-element.238 = f32[128,128]{1,0} get-tuple-element((f32[128,128]{1,0}, f32[128,128]{1,0}, f32[128,128]{1,0}) %fusion.42), index=2
  %get-tuple-element.239 = f32[3]{0} get-tuple-element((f32[3]{0}, f32[3]{0}, f32[3]{0}) %fusion.39), index=2
  %get-tuple-element.240 = f32[128,3]{1,0} get-tuple-element((f32[128,3]{1,0}, f32[128,3]{1,0}, f32[128,3]{1,0}) %fusion.36), index=2
  %get-tuple-element.241 = f32[3,48,128]{2,1,0} get-tuple-element((f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}) %fusion.31), index=2
  %get-tuple-element.242 = f32[3,48,128,128]{3,2,1,0} get-tuple-element((f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}) %fusion.27), index=2
  %get-tuple-element.243 = f32[3,16,128]{2,1,0} get-tuple-element((f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}) %fusion.7), index=2
  %get-tuple-element.244 = f32[3,16,128,128]{3,2,1,0} get-tuple-element((f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}) %fusion.3), index=2
  %get-tuple-element.245 = f32[144,27]{1,0} get-tuple-element((f32[144,27]{1,0}, f32[144,27]{1,0}, f32[144,27]{1,0}) %fusion.54), index=1
  %get-tuple-element.246 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) %fusion.45), index=4
  %get-tuple-element.247 = f32[150,128]{1,0} get-tuple-element((f32[150,128]{1,0}, f32[150,128]{1,0}, f32[150,128]{1,0}) %fusion.48), index=1
  %get-tuple-element.248 = f32[128]{0} get-tuple-element((f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, f32[128]{0}, /*index=5*/f32[128]{0}) %fusion.45), index=1
  %get-tuple-element.249 = f32[128,128]{1,0} get-tuple-element((f32[128,128]{1,0}, f32[128,128]{1,0}, f32[128,128]{1,0}) %fusion.42), index=1
  %get-tuple-element.250 = f32[3]{0} get-tuple-element((f32[3]{0}, f32[3]{0}, f32[3]{0}) %fusion.39), index=1
  %get-tuple-element.251 = f32[128,3]{1,0} get-tuple-element((f32[128,3]{1,0}, f32[128,3]{1,0}, f32[128,3]{1,0}) %fusion.36), index=1
  %get-tuple-element.252 = f32[3,48,128]{2,1,0} get-tuple-element((f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}, f32[3,48,128]{2,1,0}) %fusion.31), index=1
  %get-tuple-element.253 = f32[3,48,128,128]{3,2,1,0} get-tuple-element((f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,48,128,128]{3,2,1,0}) %fusion.27), index=1
  %get-tuple-element.254 = f32[3,16,128]{2,1,0} get-tuple-element((f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128]{2,1,0}) %fusion.7), index=1
  %get-tuple-element.255 = f32[3,16,128,128]{3,2,1,0} get-tuple-element((f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[3,16,128,128]{3,2,1,0}) %fusion.3), index=1
  %get-tuple-element.257 = u32[2]{0} get-tuple-element((u32[2]{0}, u32[2]{0}, u32[2]{0}) %fusion.2), index=0
  %get-tuple-element.258 = s32[] get-tuple-element((f32[], f32[], s32[]) %fusion.166), index=2
  %get-tuple-element.259 = f32[] get-tuple-element((f32[], f32[], s32[]) %fusion.166), index=0
  %constant_43 = f32[] constant(0.02)
  %multiply.2265 = f32[] multiply(f32[] %get-tuple-element.184, f32[] %constant_43), metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=355}
  %input_fusion_reduce = f32[] fusion(f32[4096,3]{1,0} %fusion.66), kind=kInput, calls=%input_fused_computation_reduce, metadata={op_name="jit(training_step)/jit(main)/reduce_sum[axes=(0, 1)]" source_file="/home/brent/tensorf-jax/tensorf/training.py" source_line=291}
  %fusion.1 = (f32[], f32[]) fusion(f32[] %input_fusion_reduce), kind=kLoop, calls=%fused_computation.1, metadata={op_name="jit(training_step)/jit(main)/mul" source_file="/home/brent/tensorf-jax/tensorf/utils.py" source_line=17}
  %get-tuple-element.261 = f32[] get-tuple-element((f32[], f32[]) %fusion.1), index=1
  %get-tuple-element.262 = f32[] get-tuple-element((f32[], f32[]) %fusion.1), index=0
  ROOT %tuple.91 = (f32[144,27]{1,0}, f32[128]{0}, f32[150,128]{1,0}, f32[128]{0}, f32[128,128]{1,0}, /*index=5*/f32[3]{0}, f32[128,3]{1,0}, f32[3,48,128]{2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,16,128]{2,1,0}, /*index=10*/f32[3,16,128,128]{3,2,1,0}, s32[], f32[144,27]{1,0}, f32[128]{0}, f32[150,128]{1,0}, /*index=15*/f32[128]{0}, f32[128,128]{1,0}, f32[3]{0}, f32[128,3]{1,0}, f32[3,48,128]{2,1,0}, /*index=20*/f32[3,48,128,128]{3,2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[144,27]{1,0}, f32[128]{0}, /*index=25*/f32[150,128]{1,0}, f32[128]{0}, f32[128,128]{1,0}, f32[3]{0}, f32[128,3]{1,0}, /*index=30*/f32[3,48,128]{2,1,0}, f32[3,48,128,128]{3,2,1,0}, f32[3,16,128]{2,1,0}, f32[3,16,128,128]{3,2,1,0}, f32[2,3]{1,0}, /*index=35*/u32[2]{0}, s32[], f32[], f32[], f32[], /*index=40*/f32[]) tuple(f32[144,27]{1,0} %get-tuple-element.222, f32[128]{0} %get-tuple-element.223, f32[150,128]{1,0} %get-tuple-element.224, f32[128]{0} %get-tuple-element.225, f32[128,128]{1,0} %get-tuple-element.226, /*index=5*/f32[3]{0} %get-tuple-element.227, f32[128,3]{1,0} %get-tuple-element.228, f32[3,48,128]{2,1,0} %get-tuple-element.229, f32[3,48,128,128]{3,2,1,0} %get-tuple-element.230, f32[3,16,128]{2,1,0} %get-tuple-element.231, /*index=10*/f32[3,16,128,128]{3,2,1,0} %get-tuple-element.232, s32[] %get-tuple-element.233, f32[144,27]{1,0} %get-tuple-element.234, f32[128]{0} %get-tuple-element.235, f32[150,128]{1,0} %get-tuple-element.236, /*index=15*/f32[128]{0} %get-tuple-element.237, f32[128,128]{1,0} %get-tuple-element.238, f32[3]{0} %get-tuple-element.239, f32[128,3]{1,0} %get-tuple-element.240, f32[3,48,128]{2,1,0} %get-tuple-element.241, /*index=20*/f32[3,48,128,128]{3,2,1,0} %get-tuple-element.242, f32[3,16,128]{2,1,0} %get-tuple-element.243, f32[3,16,128,128]{3,2,1,0} %get-tuple-element.244, f32[144,27]{1,0} %get-tuple-element.245, f32[128]{0} %get-tuple-element.246, /*index=25*/f32[150,128]{1,0} %get-tuple-element.247, f32[128]{0} %get-tuple-element.248, f32[128,128]{1,0} %get-tuple-element.249, f32[3]{0} %get-tuple-element.250, f32[128,3]{1,0} %get-tuple-element.251, /*index=30*/f32[3,48,128]{2,1,0} %get-tuple-element.252, f32[3,48,128,128]{3,2,1,0} %get-tuple-element.253, f32[3,16,128]{2,1,0} %get-tuple-element.254, f32[3,16,128,128]{3,2,1,0} %get-tuple-element.255, f32[2,3]{1,0} %Arg_34.35, /*index=35*/u32[2]{0} %get-tuple-element.257, s32[] %get-tuple-element.258, f32[] %get-tuple-element.259, f32[] %multiply.2265, f32[] %get-tuple-element.261, /*index=40*/f32[] %get-tuple-element.262)
}