比较文本

查找两个文本文件之间的差异

Real-time diff

Unified diff

Collapse lines

Highlight change

Syntax highlighting

工具

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

interpolate 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

98 lines

Lines
Total
Added

Words
Total
Added

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

96 lines

HloModule jit_interpolate_backward.17

%region_0.283 (Arg_0.284: f32[], Arg_1.285: f32[]) -> f32[] {

%region_0.286 (Arg_0.287: f32[], Arg_1.288: f32[]) -> f32[] {

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

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

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

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

ROOT %add.286 = f32[] add(f32[] %Arg_0.284, f32[] %Arg_1.285), metadata={op_name="jit(interpolate_backward)/jit(main)/reduce_sum[axes=(0, 1, 2)]" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=46}

ROOT %add.289 = f32[] add(f32[] %Arg_0.287, f32[] %Arg_1.288), metadata={op_name="jit(interpolate_backward)/jit(main)/reduce_sum[axes=(0, 1, 2)]" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=46}

}

%region_4.303 (Arg_0.304: f32[], Arg_1.305: f32[]) -> f32[] {

%region_4.306 (Arg_0.307: f32[], Arg_1.308: f32[]) -> f32[] {

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

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

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

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

ROOT %add.306 = f32[] add(f32[] %Arg_0.304, f32[] %Arg_1.305), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

ROOT %add.309 = f32[] add(f32[] %Arg_0.307, f32[] %Arg_1.308), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

}

%input_fused_computation_scatter (param_0.150: f32[3,48,4096,443], param_1.215: f32[3,4096,443], param_2.216: f32[3,4096,443]) -> f32[3,48,128,128] {

%input_fused_computation_scatter (param_0.154: f32[3,4096,48,443], param_1.215: f32[3,4096,443], param_2.216: f32[3,4096,443]) -> f32[3,48,128,128] {

%constant_54 = f32[] constant(0)

%broadcast.169 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_54), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%broadcast.171 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_54), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.56 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%iota.56 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%constant_387 = s32[] constant(0)

%broadcast.511 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_387), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%broadcast.512 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_387), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.216 = f32[3,4096,443]{2,1,0} parameter(2)

%floor.133 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_2.216), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%floor.133 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_2.216), metadata={op_name="jit(interpolate_backward)/jit(main)/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.76 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.133), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%convert.76 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.133), metadata={op_name="jit(interpolate_backward)/jit(main)/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_386 = s32[] constant(127)

%broadcast.510 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_386), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%broadcast.511 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_386), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.75 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.511, s32[3,4096,443]{2,1,0} %convert.76, s32[3,4096,443]{2,1,0} %broadcast.510), metadata={op_name="jit(interpolate_backward)/jit(main)/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.75 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.512, s32[3,4096,443]{2,1,0} %convert.76, s32[3,4096,443]{2,1,0} %broadcast.511), metadata={op_name="jit(interpolate_backward)/jit(main)/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.93 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.75, s32[3,4096,443]{2,1,0} %broadcast.511), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%compare.93 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.75, s32[3,4096,443]{2,1,0} %broadcast.512), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_385 = s32[] constant(128)

%broadcast.509 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_385), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%broadcast.510 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_385), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.180 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.75, s32[3,4096,443]{2,1,0} %broadcast.509), metadata={op_name="jit(interpolate_backward)/jit(main)/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.180 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.75, s32[3,4096,443]{2,1,0} %broadcast.510), metadata={op_name="jit(interpolate_backward)/jit(main)/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.93 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.93, s32[3,4096,443]{2,1,0} %add.180, s32[3,4096,443]{2,1,0} %clamp.75), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%select.93 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.93, s32[3,4096,443]{2,1,0} %add.180, s32[3,4096,443]{2,1,0} %clamp.75), metadata={op_name="jit(interpolate_backward)/jit(main)/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.107 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.93), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%bitcast.107 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.93), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%param_1.215 = f32[3,4096,443]{2,1,0} parameter(1)

%floor.132 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.215), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%floor.132 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.215), metadata={op_name="jit(interpolate_backward)/jit(main)/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.75 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.132), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%convert.75 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.132), metadata={op_name="jit(interpolate_backward)/jit(main)/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.74 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.511, s32[3,4096,443]{2,1,0} %convert.75, s32[3,4096,443]{2,1,0} %broadcast.510), metadata={op_name="jit(interpolate_backward)/jit(main)/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.74 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.512, s32[3,4096,443]{2,1,0} %convert.75, s32[3,4096,443]{2,1,0} %broadcast.511), metadata={op_name="jit(interpolate_backward)/jit(main)/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.92 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.74, s32[3,4096,443]{2,1,0} %broadcast.511), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%compare.92 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.74, s32[3,4096,443]{2,1,0} %broadcast.512), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.179 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.74, s32[3,4096,443]{2,1,0} %broadcast.509), metadata={op_name="jit(interpolate_backward)/jit(main)/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.179 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.74, s32[3,4096,443]{2,1,0} %broadcast.510), metadata={op_name="jit(interpolate_backward)/jit(main)/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.92 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.92, s32[3,4096,443]{2,1,0} %add.179, s32[3,4096,443]{2,1,0} %clamp.74), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%select.92 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.92, s32[3,4096,443]{2,1,0} %add.179, s32[3,4096,443]{2,1,0} %clamp.74), metadata={op_name="jit(interpolate_backward)/jit(main)/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.106 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.92), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%bitcast.106 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.92), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%concatenate.58 = s32[3,4096,443,3]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.56, s32[3,4096,443,1]{2,1,0,3} %bitcast.107, s32[3,4096,443,1]{2,1,0,3} %bitcast.106), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%concatenate.58 = s32[3,4096,443,3]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.56, s32[3,4096,443,1]{2,1,0,3} %bitcast.107, s32[3,4096,443,1]{2,1,0,3} %bitcast.106), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.8 = s32[3,4096,443,3]{3,2,1,0} copy(s32[3,4096,443,3]{2,1,0,3} %concatenate.58), metadata={op_name="jit(interpolate_backward)/jit(main)/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.8 = s32[3,4096,443,3]{3,2,1,0} copy(s32[3,4096,443,3]{2,1,0,3} %concatenate.58), metadata={op_name="jit(interpolate_backward)/jit(main)/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_364 = f32[] constant(1)

%broadcast.487 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_364), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.488 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_364), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.97 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_2.216, f32[3,4096,443]{2,1,0} %floor.133), metadata={op_name="jit(interpolate_backward)/jit(main)/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.97 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_2.216, f32[3,4096,443]{2,1,0} %floor.133), metadata={op_name="jit(interpolate_backward)/jit(main)/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.96 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.487, f32[3,4096,443]{2,1,0} %subtract.97), metadata={op_name="jit(interpolate_backward)/jit(main)/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.96 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.488, f32[3,4096,443]{2,1,0} %subtract.97), metadata={op_name="jit(interpolate_backward)/jit(main)/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.95 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.215, f32[3,4096,443]{2,1,0} %floor.132), metadata={op_name="jit(interpolate_backward)/jit(main)/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.95 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.215, f32[3,4096,443]{2,1,0} %floor.132), metadata={op_name="jit(interpolate_backward)/jit(main)/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.94 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.487, f32[3,4096,443]{2,1,0} %subtract.95), metadata={op_name="jit(interpolate_backward)/jit(main)/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.94 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.488, f32[3,4096,443]{2,1,0} %subtract.95), metadata={op_name="jit(interpolate_backward)/jit(main)/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.93 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.96, f32[3,4096,443]{2,1,0} %subtract.94), metadata={op_name="jit(interpolate_backward)/jit(main)/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.94 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.96, f32[3,4096,443]{2,1,0} %subtract.94), metadata={op_name="jit(interpolate_backward)/jit(main)/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.486 = f32[3,48,4096,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.93), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.487 = f32[3,4096,48,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.94), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.150 = f32[3,48,4096,443]{3,2,1,0} parameter(0)

%param_0.154 = f32[3,4096,48,443]{3,2,1,0} parameter(0)

%multiply.40 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %broadcast.486, f32[3,48,4096,443]{3,2,1,0} %param_0.150), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%multiply.41 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %broadcast.487, f32[3,4096,48,443]{3,2,1,0} %param_0.154), metadata={op_name="jit(interpolate_backward)/jit(main)/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.12 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.169, s32[3,4096,443,3]{3,2,1,0} %copy.8, f32[3,48,4096,443]{3,2,1,0} %multiply.40), 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_4.303, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

ROOT %scatter.12 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.171, s32[3,4096,443,3]{3,2,1,0} %copy.8, f32[3,4096,48,443]{3,2,1,0} %multiply.41), 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_4.306, metadata={op_name="jit(interpolate_backward)/jit(main)/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_3.299 (Arg_0.300: f32[], Arg_1.301: f32[]) -> f32[] {

%region_3.302 (Arg_0.303: f32[], Arg_1.304: f32[]) -> f32[] {

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

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

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

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

ROOT %add.302 = f32[] add(f32[] %Arg_0.300, f32[] %Arg_1.301), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

ROOT %add.305 = f32[] add(f32[] %Arg_0.303, f32[] %Arg_1.304), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

}

%param_1.207 = f32[3,48,128,128]{3,2,1,0} parameter(1)

%param_2.209 = f32[3,48,128,128]{3,2,1,0} parameter(2)

%add.50 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %param_1.207, f32[3,48,128,128]{3,2,1,0} %param_2.209), metadata={op_name="jit(interpolate_backward)/jit(main)/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.50 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %param_1.207, f32[3,48,128,128]{3,2,1,0} %param_2.209), metadata={op_name="jit(interpolate_backward)/jit(main)/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.52 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%iota.52 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%constant_354 = s32[] constant(0)

%broadcast.479 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_354), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%broadcast.480 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_354), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.100 = f32[3,4096,443]{2,1,0} parameter(4)

%floor.115 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_4.100), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%floor.115 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_4.100), metadata={op_name="jit(interpolate_backward)/jit(main)/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.68 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.115), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%convert.68 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.115), metadata={op_name="jit(interpolate_backward)/jit(main)/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_352 = s32[] constant(127)

%broadcast.478 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_352), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%broadcast.479 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_352), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.479, s32[3,4096,443]{2,1,0} %convert.68, s32[3,4096,443]{2,1,0} %broadcast.478), metadata={op_name="jit(interpolate_backward)/jit(main)/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.67 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.480, s32[3,4096,443]{2,1,0} %convert.68, s32[3,4096,443]{2,1,0} %broadcast.479), metadata={op_name="jit(interpolate_backward)/jit(main)/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.85 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.67, s32[3,4096,443]{2,1,0} %broadcast.479), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%compare.85 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.67, s32[3,4096,443]{2,1,0} %broadcast.480), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_351 = s32[] constant(128)

%broadcast.477 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_351), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%broadcast.478 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_351), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.170 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.67, s32[3,4096,443]{2,1,0} %broadcast.477), metadata={op_name="jit(interpolate_backward)/jit(main)/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.170 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.67, s32[3,4096,443]{2,1,0} %broadcast.478), metadata={op_name="jit(interpolate_backward)/jit(main)/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.85 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.85, s32[3,4096,443]{2,1,0} %add.170, s32[3,4096,443]{2,1,0} %clamp.67), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%select.85 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.85, s32[3,4096,443]{2,1,0} %add.170, s32[3,4096,443]{2,1,0} %clamp.67), metadata={op_name="jit(interpolate_backward)/jit(main)/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.99 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.85), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%bitcast.99 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.85), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%param_3.173 = f32[3,4096,443]{2,1,0} parameter(3)

%floor.114 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_3.173), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%floor.114 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_3.173), metadata={op_name="jit(interpolate_backward)/jit(main)/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.67 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.114), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%convert.67 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.114), metadata={op_name="jit(interpolate_backward)/jit(main)/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_353 = s32[] constant(1)

%broadcast.475 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_353), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%broadcast.476 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_353), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.169 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.67, s32[3,4096,443]{2,1,0} %broadcast.475), metadata={op_name="jit(interpolate_backward)/jit(main)/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.169 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.67, s32[3,4096,443]{2,1,0} %broadcast.476), metadata={op_name="jit(interpolate_backward)/jit(main)/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.66 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.479, s32[3,4096,443]{2,1,0} %add.169, s32[3,4096,443]{2,1,0} %broadcast.478), metadata={op_name="jit(interpolate_backward)/jit(main)/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.66 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.480, s32[3,4096,443]{2,1,0} %add.169, s32[3,4096,443]{2,1,0} %broadcast.479), metadata={op_name="jit(interpolate_backward)/jit(main)/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.84 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.66, s32[3,4096,443]{2,1,0} %broadcast.479), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%compare.84 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.66, s32[3,4096,443]{2,1,0} %broadcast.480), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.168 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.66, s32[3,4096,443]{2,1,0} %broadcast.477), metadata={op_name="jit(interpolate_backward)/jit(main)/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.168 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.66, s32[3,4096,443]{2,1,0} %broadcast.478), metadata={op_name="jit(interpolate_backward)/jit(main)/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.84 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.84, s32[3,4096,443]{2,1,0} %add.168, s32[3,4096,443]{2,1,0} %clamp.66), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%select.84 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.84, s32[3,4096,443]{2,1,0} %add.168, s32[3,4096,443]{2,1,0} %clamp.66), metadata={op_name="jit(interpolate_backward)/jit(main)/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.98 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.84), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%bitcast.98 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.84), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}

%concatenate.54 = s32[3,4096,443,3]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.52, s32[3,4096,443,1]{3,2,1,0} %bitcast.99, s32[3,4096,443,1]{3,2,1,0} %bitcast.98), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}

%concatenate.54 = s32[3,4096,443,3]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.52, s32[3,4096,443,1]{3,2,1,0} %bitcast.99, s32[3,4096,443,1]{3,2,1,0} %bitcast.98), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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_312 = f32[] constant(1)

%broadcast.439 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_312), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.440 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_312), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.81 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_4.100, f32[3,4096,443]{2,1,0} %floor.115), metadata={op_name="jit(interpolate_backward)/jit(main)/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.81 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_4.100, f32[3,4096,443]{2,1,0} %floor.115), metadata={op_name="jit(interpolate_backward)/jit(main)/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.80 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.439, f32[3,4096,443]{2,1,0} %subtract.81), metadata={op_name="jit(interpolate_backward)/jit(main)/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.80 = f32[3,4096,443]{2,1,0}

%subtract.79 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_3.173, f32[3,4096,443]{2,1,0} %floor.114), metadata={op_name="jit(interpolate_backward)/jit(main)/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}

已保存差异

原始文本

打开文件

HloModule jit_interpolate_backward.17

%region_0.283 (Arg_0.284: f32[], Arg_1.285: f32[]) -> f32[] {
  %Arg_0.284 = f32[] parameter(0)
  %Arg_1.285 = f32[] parameter(1)
  ROOT %add.286 = f32[] add(f32[] %Arg_0.284, f32[] %Arg_1.285), metadata={op_name="jit(interpolate_backward)/jit(main)/reduce_sum[axes=(0, 1, 2)]" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=46}
}

%region_4.303 (Arg_0.304: f32[], Arg_1.305: f32[]) -> f32[] {
  %Arg_0.304 = f32[] parameter(0)
  %Arg_1.305 = f32[] parameter(1)
  ROOT %add.306 = f32[] add(f32[] %Arg_0.304, f32[] %Arg_1.305), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%input_fused_computation_scatter (param_0.150: f32[3,48,4096,443], param_1.215: f32[3,4096,443], param_2.216: f32[3,4096,443]) -> f32[3,48,128,128] {
  %constant_54 = f32[] constant(0)
  %broadcast.169 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_54), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.56 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_387 = s32[] constant(0)
  %broadcast.511 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_387), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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_2.216 = f32[3,4096,443]{2,1,0} parameter(2)
  %floor.133 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_2.216), metadata={op_name="jit(interpolate_backward)/jit(main)/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.76 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.133), metadata={op_name="jit(interpolate_backward)/jit(main)/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_386 = s32[] constant(127)
  %broadcast.510 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_386), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.75 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.511, s32[3,4096,443]{2,1,0} %convert.76, s32[3,4096,443]{2,1,0} %broadcast.510), metadata={op_name="jit(interpolate_backward)/jit(main)/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.93 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.75, s32[3,4096,443]{2,1,0} %broadcast.511), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_385 = s32[] constant(128)
  %broadcast.509 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_385), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.180 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.75, s32[3,4096,443]{2,1,0} %broadcast.509), metadata={op_name="jit(interpolate_backward)/jit(main)/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.93 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.93, s32[3,4096,443]{2,1,0} %add.180, s32[3,4096,443]{2,1,0} %clamp.75), metadata={op_name="jit(interpolate_backward)/jit(main)/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.107 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.93), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.215 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.132 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.215), metadata={op_name="jit(interpolate_backward)/jit(main)/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.75 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.132), metadata={op_name="jit(interpolate_backward)/jit(main)/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.74 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.511, s32[3,4096,443]{2,1,0} %convert.75, s32[3,4096,443]{2,1,0} %broadcast.510), metadata={op_name="jit(interpolate_backward)/jit(main)/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.92 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.74, s32[3,4096,443]{2,1,0} %broadcast.511), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.179 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.74, s32[3,4096,443]{2,1,0} %broadcast.509), metadata={op_name="jit(interpolate_backward)/jit(main)/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.92 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.92, s32[3,4096,443]{2,1,0} %add.179, s32[3,4096,443]{2,1,0} %clamp.74), metadata={op_name="jit(interpolate_backward)/jit(main)/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.106 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.92), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.58 = s32[3,4096,443,3]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.56, s32[3,4096,443,1]{2,1,0,3} %bitcast.107, s32[3,4096,443,1]{2,1,0,3} %bitcast.106), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.8 = s32[3,4096,443,3]{3,2,1,0} copy(s32[3,4096,443,3]{2,1,0,3} %concatenate.58), metadata={op_name="jit(interpolate_backward)/jit(main)/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_364 = f32[] constant(1)
  %broadcast.487 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_364), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.97 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_2.216, f32[3,4096,443]{2,1,0} %floor.133), metadata={op_name="jit(interpolate_backward)/jit(main)/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.96 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.487, f32[3,4096,443]{2,1,0} %subtract.97), metadata={op_name="jit(interpolate_backward)/jit(main)/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.95 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.215, f32[3,4096,443]{2,1,0} %floor.132), metadata={op_name="jit(interpolate_backward)/jit(main)/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.94 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.487, f32[3,4096,443]{2,1,0} %subtract.95), metadata={op_name="jit(interpolate_backward)/jit(main)/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.93 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.96, f32[3,4096,443]{2,1,0} %subtract.94), metadata={op_name="jit(interpolate_backward)/jit(main)/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.486 = f32[3,48,4096,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.93), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.150 = f32[3,48,4096,443]{3,2,1,0} parameter(0)
  %multiply.40 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %broadcast.486, f32[3,48,4096,443]{3,2,1,0} %param_0.150), metadata={op_name="jit(interpolate_backward)/jit(main)/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.12 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.169, s32[3,4096,443,3]{3,2,1,0} %copy.8, f32[3,48,4096,443]{3,2,1,0} %multiply.40), 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_4.303, metadata={op_name="jit(interpolate_backward)/jit(main)/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_3.299 (Arg_0.300: f32[], Arg_1.301: f32[]) -> f32[] {
  %Arg_0.300 = f32[] parameter(0)
  %Arg_1.301 = f32[] parameter(1)
  ROOT %add.302 = f32[] add(f32[] %Arg_0.300, f32[] %Arg_1.301), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%input_fused_computation_scatter.1 (param_0.147: f32[3,48,4096,443], param_1.207: f32[3,48,128,128], param_2.209: f32[3,48,128,128], param_3.173: f32[3,4096,443], param_4.100: f32[3,4096,443]) -> f32[3,48,128,128] {
  %param_1.207 = f32[3,48,128,128]{3,2,1,0} parameter(1)
  %param_2.209 = f32[3,48,128,128]{3,2,1,0} parameter(2)
  %add.50 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %param_1.207, f32[3,48,128,128]{3,2,1,0} %param_2.209), metadata={op_name="jit(interpolate_backward)/jit(main)/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.52 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_354 = s32[] constant(0)
  %broadcast.479 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_354), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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_4.100 = f32[3,4096,443]{2,1,0} parameter(4)
  %floor.115 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_4.100), metadata={op_name="jit(interpolate_backward)/jit(main)/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.68 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.115), metadata={op_name="jit(interpolate_backward)/jit(main)/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_352 = s32[] constant(127)
  %broadcast.478 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_352), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.67 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.479, s32[3,4096,443]{2,1,0} %convert.68, s32[3,4096,443]{2,1,0} %broadcast.478), metadata={op_name="jit(interpolate_backward)/jit(main)/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.85 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.67, s32[3,4096,443]{2,1,0} %broadcast.479), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_351 = s32[] constant(128)
  %broadcast.477 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_351), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.170 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.67, s32[3,4096,443]{2,1,0} %broadcast.477), metadata={op_name="jit(interpolate_backward)/jit(main)/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.85 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.85, s32[3,4096,443]{2,1,0} %add.170, s32[3,4096,443]{2,1,0} %clamp.67), metadata={op_name="jit(interpolate_backward)/jit(main)/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.99 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.85), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_3.173 = f32[3,4096,443]{2,1,0} parameter(3)
  %floor.114 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_3.173), metadata={op_name="jit(interpolate_backward)/jit(main)/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.67 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.114), metadata={op_name="jit(interpolate_backward)/jit(main)/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_353 = s32[] constant(1)
  %broadcast.475 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_353), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.169 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.67, s32[3,4096,443]{2,1,0} %broadcast.475), metadata={op_name="jit(interpolate_backward)/jit(main)/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.66 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.479, s32[3,4096,443]{2,1,0} %add.169, s32[3,4096,443]{2,1,0} %broadcast.478), metadata={op_name="jit(interpolate_backward)/jit(main)/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.84 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.66, s32[3,4096,443]{2,1,0} %broadcast.479), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.168 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.66, s32[3,4096,443]{2,1,0} %broadcast.477), metadata={op_name="jit(interpolate_backward)/jit(main)/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.84 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.84, s32[3,4096,443]{2,1,0} %add.168, s32[3,4096,443]{2,1,0} %clamp.66), metadata={op_name="jit(interpolate_backward)/jit(main)/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.98 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.84), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.54 = s32[3,4096,443,3]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.52, s32[3,4096,443,1]{3,2,1,0} %bitcast.99, s32[3,4096,443,1]{3,2,1,0} %bitcast.98), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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_312 = f32[] constant(1)
  %broadcast.439 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_312), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.81 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_4.100, f32[3,4096,443]{2,1,0} %floor.115), metadata={op_name="jit(interpolate_backward)/jit(main)/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.80 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.439, f32[3,4096,443]{2,1,0} %subtract.81), metadata={op_name="jit(interpolate_backward)/jit(main)/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.79 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_3.173, f32[3,4096,443]{2,1,0} %floor.114), metadata={op_name="jit(interpolate_backward)/jit(main)/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.89 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.80, f32[3,4096,443]{2,1,0} %subtract.79), metadata={op_name="jit(interpolate_backward)/jit(main)/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.438 = f32[3,48,4096,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.89), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.147 = f32[3,48,4096,443]{3,2,1,0} parameter(0)
  %multiply.41 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %broadcast.438, f32[3,48,4096,443]{3,2,1,0} %param_0.147), metadata={op_name="jit(interpolate_backward)/jit(main)/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.13 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %add.50, s32[3,4096,443,3]{3,2,1,0} %concatenate.54, f32[3,48,4096,443]{3,2,1,0} %multiply.41), 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_3.299, metadata={op_name="jit(interpolate_backward)/jit(main)/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_2.295 (Arg_0.296: f32[], Arg_1.297: f32[]) -> f32[] {
  %Arg_0.296 = f32[] parameter(0)
  %Arg_1.297 = f32[] parameter(1)
  ROOT %add.298 = f32[] add(f32[] %Arg_0.296, f32[] %Arg_1.297), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%input_fused_computation_scatter.2 (param_0.132: f32[3,48,4096,443], param_1.191: f32[3,4096,443], param_2.191: f32[3,4096,443]) -> f32[3,48,128,128] {
  %constant_55 = f32[] constant(0)
  %broadcast.171 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_55), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.47 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_302 = s32[] constant(0)
  %broadcast.431 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_302), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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_2.191 = f32[3,4096,443]{2,1,0} parameter(2)
  %floor.95 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_2.191), metadata={op_name="jit(interpolate_backward)/jit(main)/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.56 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.95), metadata={op_name="jit(interpolate_backward)/jit(main)/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_305 = s32[] constant(1)
  %broadcast.430 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_305), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.152 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.56, s32[3,4096,443]{2,1,0} %broadcast.430), metadata={op_name="jit(interpolate_backward)/jit(main)/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_301 = s32[] constant(127)
  %broadcast.429 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_301), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.55 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.431, s32[3,4096,443]{2,1,0} %add.152, s32[3,4096,443]{2,1,0} %broadcast.429), metadata={op_name="jit(interpolate_backward)/jit(main)/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.73 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.55, s32[3,4096,443]{2,1,0} %broadcast.431), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_300 = s32[] constant(128)
  %broadcast.428 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_300), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.151 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.55, s32[3,4096,443]{2,1,0} %broadcast.428), metadata={op_name="jit(interpolate_backward)/jit(main)/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.73 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.73, s32[3,4096,443]{2,1,0} %add.151, s32[3,4096,443]{2,1,0} %clamp.55), metadata={op_name="jit(interpolate_backward)/jit(main)/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.91 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.73), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.191 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.94 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.191), metadata={op_name="jit(interpolate_backward)/jit(main)/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.55 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.94), metadata={op_name="jit(interpolate_backward)/jit(main)/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.54 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.431, s32[3,4096,443]{2,1,0} %convert.55, s32[3,4096,443]{2,1,0} %broadcast.429), metadata={op_name="jit(interpolate_backward)/jit(main)/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.72 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.54, s32[3,4096,443]{2,1,0} %broadcast.431), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.150 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.54, s32[3,4096,443]{2,1,0} %broadcast.428), metadata={op_name="jit(interpolate_backward)/jit(main)/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.72 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.72, s32[3,4096,443]{2,1,0} %add.150, s32[3,4096,443]{2,1,0} %clamp.54), metadata={op_name="jit(interpolate_backward)/jit(main)/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.90 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.72), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.49 = s32[3,4096,443,3]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.47, s32[3,4096,443,1]{2,1,0,3} %bitcast.91, s32[3,4096,443,1]{2,1,0,3} %bitcast.90), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.9 = s32[3,4096,443,3]{3,2,1,0} copy(s32[3,4096,443,3]{2,1,0,3} %concatenate.49), metadata={op_name="jit(interpolate_backward)/jit(main)/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.69 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_2.191, f32[3,4096,443]{2,1,0} %floor.95), metadata={op_name="jit(interpolate_backward)/jit(main)/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_262 = f32[] constant(1)
  %broadcast.391 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_262), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.68 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.191, f32[3,4096,443]{2,1,0} %floor.94), metadata={op_name="jit(interpolate_backward)/jit(main)/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.67 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.391, f32[3,4096,443]{2,1,0} %subtract.68), metadata={op_name="jit(interpolate_backward)/jit(main)/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.85 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.69, f32[3,4096,443]{2,1,0} %subtract.67), metadata={op_name="jit(interpolate_backward)/jit(main)/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.390 = f32[3,48,4096,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.85), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.132 = f32[3,48,4096,443]{3,2,1,0} parameter(0)
  %multiply.42 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %broadcast.390, f32[3,48,4096,443]{3,2,1,0} %param_0.132), metadata={op_name="jit(interpolate_backward)/jit(main)/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.14 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.171, s32[3,4096,443,3]{3,2,1,0} %copy.9, f32[3,48,4096,443]{3,2,1,0} %multiply.42), 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_2.295, metadata={op_name="jit(interpolate_backward)/jit(main)/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_1.291 (Arg_0.292: f32[], Arg_1.293: f32[]) -> f32[] {
  %Arg_0.292 = f32[] parameter(0)
  %Arg_1.293 = f32[] parameter(1)
  ROOT %add.294 = f32[] add(f32[] %Arg_0.292, f32[] %Arg_1.293), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%input_fused_computation_scatter.3 (param_0.129: f32[3,48,4096,443], param_1.173: f32[3,4096,443], param_2.171: f32[3,4096,443]) -> f32[3,48,128,128] {
  %constant_56 = f32[] constant(0)
  %broadcast.172 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_56), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.43 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_252 = s32[] constant(0)
  %broadcast.383 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_252), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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_2.171 = f32[3,4096,443]{2,1,0} parameter(2)
  %floor.75 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_2.171), metadata={op_name="jit(interpolate_backward)/jit(main)/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.44 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.75), metadata={op_name="jit(interpolate_backward)/jit(main)/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_251 = s32[] constant(1)
  %broadcast.382 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_251), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.134 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.44, s32[3,4096,443]{2,1,0} %broadcast.382), metadata={op_name="jit(interpolate_backward)/jit(main)/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_250 = s32[] constant(127)
  %broadcast.381 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_250), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.43 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.383, s32[3,4096,443]{2,1,0} %add.134, s32[3,4096,443]{2,1,0} %broadcast.381), metadata={op_name="jit(interpolate_backward)/jit(main)/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.61 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.43, s32[3,4096,443]{2,1,0} %broadcast.383), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_249 = s32[] constant(128)
  %broadcast.380 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_249), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.132 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.43, s32[3,4096,443]{2,1,0} %broadcast.380), metadata={op_name="jit(interpolate_backward)/jit(main)/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.61 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.61, s32[3,4096,443]{2,1,0} %add.132, s32[3,4096,443]{2,1,0} %clamp.43), metadata={op_name="jit(interpolate_backward)/jit(main)/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.79 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.61), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.173 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.74 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.173), metadata={op_name="jit(interpolate_backward)/jit(main)/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.43 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.74), metadata={op_name="jit(interpolate_backward)/jit(main)/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.131 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.43, s32[3,4096,443]{2,1,0} %broadcast.382), metadata={op_name="jit(interpolate_backward)/jit(main)/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.42 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.383, s32[3,4096,443]{2,1,0} %add.131, s32[3,4096,443]{2,1,0} %broadcast.381), metadata={op_name="jit(interpolate_backward)/jit(main)/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.60 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.42, s32[3,4096,443]{2,1,0} %broadcast.383), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.130 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.42, s32[3,4096,443]{2,1,0} %broadcast.380), metadata={op_name="jit(interpolate_backward)/jit(main)/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.59 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.60, s32[3,4096,443]{2,1,0} %add.130, s32[3,4096,443]{2,1,0} %clamp.42), metadata={op_name="jit(interpolate_backward)/jit(main)/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.78 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.59), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.45 = s32[3,4096,443,3]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.43, s32[3,4096,443,1]{3,2,1,0} %bitcast.79, s32[3,4096,443,1]{3,2,1,0} %bitcast.78), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.57 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_2.171, f32[3,4096,443]{2,1,0} %floor.75), metadata={op_name="jit(interpolate_backward)/jit(main)/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.56 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.173, f32[3,4096,443]{2,1,0} %floor.74), metadata={op_name="jit(interpolate_backward)/jit(main)/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.81 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.57, f32[3,4096,443]{2,1,0} %subtract.56), metadata={op_name="jit(interpolate_backward)/jit(main)/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.305 = f32[3,48,4096,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.81), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.129 = f32[3,48,4096,443]{3,2,1,0} parameter(0)
  %multiply.43 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %broadcast.305, f32[3,48,4096,443]{3,2,1,0} %param_0.129), metadata={op_name="jit(interpolate_backward)/jit(main)/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.15 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.172, s32[3,4096,443,3]{3,2,1,0} %concatenate.45, f32[3,48,4096,443]{3,2,1,0} %multiply.43), 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_1.291, metadata={op_name="jit(interpolate_backward)/jit(main)/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_6.346 (Arg_0.347: f32[], Arg_1.348: f32[]) -> f32[] {
  %Arg_0.347 = f32[] parameter(0)
  %Arg_1.348 = f32[] parameter(1)
  ROOT %add.349 = f32[] add(f32[] %Arg_0.347, f32[] %Arg_1.348), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%input_fused_computation_scatter.4 (param_0.158: f32[3,48,4096,443], param_1.221: f32[3,4096,443], param_2.222: f32[3,48,4096,443], param_3.192: f32[]) -> f32[3,48,128] {
  %constant_57 = f32[] constant(0)
  %broadcast.173 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_57), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.31 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_145 = s32[] constant(0)
  %broadcast.260 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_145), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.221 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.34 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.221), metadata={op_name="jit(interpolate_backward)/jit(main)/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.19 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.34), metadata={op_name="jit(interpolate_backward)/jit(main)/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_144 = s32[] constant(127)
  %broadcast.259 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_144), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.19 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.260, s32[3,4096,443]{2,1,0} %convert.19, s32[3,4096,443]{2,1,0} %broadcast.259), metadata={op_name="jit(interpolate_backward)/jit(main)/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.35 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.19, s32[3,4096,443]{2,1,0} %broadcast.260), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_143 = s32[] constant(128)
  %broadcast.258 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_143), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.84 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.19, s32[3,4096,443]{2,1,0} %broadcast.258), metadata={op_name="jit(interpolate_backward)/jit(main)/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.35 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.35, s32[3,4096,443]{2,1,0} %add.84, s32[3,4096,443]{2,1,0} %clamp.19), metadata={op_name="jit(interpolate_backward)/jit(main)/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.59 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.35), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.33 = s32[3,4096,443,2]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.31, s32[3,4096,443,1]{3,2,1,0} %bitcast.59), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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_131 = f32[] constant(1)
  %broadcast.248 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_131), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.35 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.221, f32[3,4096,443]{2,1,0} %floor.34), metadata={op_name="jit(interpolate_backward)/jit(main)/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.34 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.248, f32[3,4096,443]{2,1,0} %subtract.35), metadata={op_name="jit(interpolate_backward)/jit(main)/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.247 = f32[3,48,4096,443]{2,1,0,3} broadcast(f32[3,4096,443]{2,1,0} %subtract.34), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.10 = f32[3,48,4096,443]{3,2,1,0} copy(f32[3,48,4096,443]{2,1,0,3} %broadcast.247), metadata={op_name="jit(interpolate_backward)/jit(main)/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_3.192 = f32[] parameter(3)
  %divide.7 = f32[] divide(f32[] %constant_131, f32[] %param_3.192), metadata={op_name="jit(interpolate_backward)/jit(main)/div" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=45}
  %broadcast.517 = f32[3,48,4096,443]{3,2,1,0} broadcast(f32[] %divide.7), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/broadcast_in_dim[shape=(144, 4096, 443) broadcast_dimensions=()]" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=46}
  %param_2.222 = f32[3,48,4096,443]{3,2,1,0} parameter(2)
  %multiply.99 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %broadcast.517, f32[3,48,4096,443]{3,2,1,0} %param_2.222), metadata={op_name="jit(interpolate_backward)/jit(main)/mul" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=47}
  %param_0.158 = f32[3,48,4096,443]{3,2,1,0} parameter(0)
  %multiply.60 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %multiply.99, f32[3,48,4096,443]{3,2,1,0} %param_0.158), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %multiply.44 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %copy.10, f32[3,48,4096,443]{3,2,1,0} %multiply.60), metadata={op_name="jit(interpolate_backward)/jit(main)/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.16 = f32[3,48,128]{2,1,0} scatter(f32[3,48,128]{2,1,0} %broadcast.173, s32[3,4096,443,2]{3,2,1,0} %concatenate.33, f32[3,48,4096,443]{3,2,1,0} %multiply.44), update_window_dims={1}, inserted_window_dims={0,2}, scatter_dims_to_operand_dims={0,2}, index_vector_dim=3, to_apply=%region_6.346, metadata={op_name="jit(interpolate_backward)/jit(main)/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_5.342 (Arg_0.343: f32[], Arg_1.344: f32[]) -> f32[] {
  %Arg_0.343 = f32[] parameter(0)
  %Arg_1.344 = f32[] parameter(1)
  ROOT %add.345 = f32[] add(f32[] %Arg_0.343, f32[] %Arg_1.344), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%input_fused_computation_scatter.5 (param_0.156: f32[3,48,4096,443], param_1.219: f32[3,4096,443], param_2.220: f32[3,48,4096,443], param_3.188: f32[]) -> f32[3,48,128] {
  %constant_58 = f32[] constant(0)
  %broadcast.175 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_58), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.27 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_125 = s32[] constant(0)
  %broadcast.240 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_125), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.219 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.26 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.219), metadata={op_name="jit(interpolate_backward)/jit(main)/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.15 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.26), metadata={op_name="jit(interpolate_backward)/jit(main)/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_124 = s32[] constant(1)
  %broadcast.239 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_124), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.80 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.15, s32[3,4096,443]{2,1,0} %broadcast.239), metadata={op_name="jit(interpolate_backward)/jit(main)/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_123 = s32[] constant(127)
  %broadcast.238 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_123), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.15 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.240, s32[3,4096,443]{2,1,0} %add.80, s32[3,4096,443]{2,1,0} %broadcast.238), metadata={op_name="jit(interpolate_backward)/jit(main)/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.31 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.15, s32[3,4096,443]{2,1,0} %broadcast.240), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_122 = s32[] constant(128)
  %broadcast.237 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_122), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.79 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.15, s32[3,4096,443]{2,1,0} %broadcast.237), metadata={op_name="jit(interpolate_backward)/jit(main)/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.31 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.31, s32[3,4096,443]{2,1,0} %add.79, s32[3,4096,443]{2,1,0} %clamp.15), metadata={op_name="jit(interpolate_backward)/jit(main)/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.55 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.31), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.29 = s32[3,4096,443,2]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.27, s32[3,4096,443,1]{2,1,0,3} %bitcast.55), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.13 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.219, f32[3,4096,443]{2,1,0} %floor.26), metadata={op_name="jit(interpolate_backward)/jit(main)/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.174 = f32[3,48,4096,443]{2,1,0,3} broadcast(f32[3,4096,443]{2,1,0} %subtract.13), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %constant_396 = f32[] constant(1)
  %param_3.188 = f32[] parameter(3)
  %divide.5 = f32[] divide(f32[] %constant_396, f32[] %param_3.188), metadata={op_name="jit(interpolate_backward)/jit(main)/div" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=45}
  %broadcast.515 = f32[3,48,4096,443]{3,2,1,0} broadcast(f32[] %divide.5), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/broadcast_in_dim[shape=(144, 4096, 443) broadcast_dimensions=()]" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=46}
  %param_2.220 = f32[3,48,4096,443]{3,2,1,0} parameter(2)
  %multiply.97 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %broadcast.515, f32[3,48,4096,443]{3,2,1,0} %param_2.220), metadata={op_name="jit(interpolate_backward)/jit(main)/mul" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=47}
  %param_0.156 = f32[3,48,4096,443]{3,2,1,0} parameter(0)
  %multiply.61 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %multiply.97, f32[3,48,4096,443]{3,2,1,0} %param_0.156), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %copy.11 = f32[3,48,4096,443]{2,1,0,3} copy(f32[3,48,4096,443]{3,2,1,0} %multiply.61), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %multiply.45 = f32[3,48,4096,443]{2,1,0,3} multiply(f32[3,48,4096,443]{2,1,0,3} %broadcast.174, f32[3,48,4096,443]{2,1,0,3} %copy.11), metadata={op_name="jit(interpolate_backward)/jit(main)/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.17 = f32[3,48,128]{2,1,0} scatter(f32[3,48,128]{2,1,0} %broadcast.175, s32[3,4096,443,2]{2,1,0,3} %concatenate.29, f32[3,48,4096,443]{2,1,0,3} %multiply.45), update_window_dims={1}, inserted_window_dims={0,2}, scatter_dims_to_operand_dims={0,2}, index_vector_dim=3, to_apply=%region_5.342, metadata={op_name="jit(interpolate_backward)/jit(main)/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_fused_computation_reduce (param_0.68: f32[3,48,4096,443]) -> f32[16165] {
  %param_0.68 = f32[3,48,4096,443]{3,2,1,0} parameter(0)
  %bitcast.19 = f32[261292032]{0} bitcast(f32[3,48,4096,443]{3,2,1,0} %param_0.68)
  %constant_60 = f32[] constant(0)
  %pad.1 = f32[261307225]{0} pad(f32[261292032]{0} %bitcast.19, f32[] %constant_60), padding=0_15193
  %bitcast.18 = f32[16165,16165]{1,0} bitcast(f32[261307225]{0} %pad.1)
  ROOT %reduce.4 = f32[16165]{0} reduce(f32[16165,16165]{1,0} %bitcast.18, f32[] %constant_60), dimensions={1}, to_apply=%region_0.283
}

%fused_computation.1 (param_0.159: f32[3,48,128], param_1.222: f32[3,4096,443], param_2.224: f32[3,48,128,128], param_3.196: f32[3,4096,443], param_4.128: f32[3,4096,443]) -> (f32[3,48,4096,443], f32[3,48,4096,443]) {
  %constant_162 = f32[] constant(1)
  %broadcast.280 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_162), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.222 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.43 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.222), metadata={op_name="jit(interpolate_backward)/jit(main)/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.43 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.222, f32[3,4096,443]{2,1,0} %floor.43), metadata={op_name="jit(interpolate_backward)/jit(main)/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.42 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.280, f32[3,4096,443]{2,1,0} %subtract.43), metadata={op_name="jit(interpolate_backward)/jit(main)/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.279 = f32[3,48,4096,443]{2,1,0,3} broadcast(f32[3,4096,443]{2,1,0} %subtract.42), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.159 = f32[3,48,128]{2,1,0} parameter(0)
  %iota.35 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_158 = s32[] constant(0)
  %broadcast.278 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_158), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.23 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.43), metadata={op_name="jit(interpolate_backward)/jit(main)/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_156 = s32[] constant(127)
  %broadcast.277 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_156), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.23 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.278, s32[3,4096,443]{2,1,0} %convert.23, s32[3,4096,443]{2,1,0} %broadcast.277), metadata={op_name="jit(interpolate_backward)/jit(main)/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.39 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.23, s32[3,4096,443]{2,1,0} %broadcast.278), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_155 = s32[] constant(128)
  %broadcast.276 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_155), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.92 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.23, s32[3,4096,443]{2,1,0} %broadcast.276), metadata={op_name="jit(interpolate_backward)/jit(main)/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.39 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.39, s32[3,4096,443]{2,1,0} %add.92, s32[3,4096,443]{2,1,0} %clamp.23), metadata={op_name="jit(interpolate_backward)/jit(main)/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.63 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.39), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.37 = s32[3,4096,443,2]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.35, s32[3,4096,443,1]{3,2,1,0} %bitcast.63), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.37 = f32[3,48,4096,443]{2,1,0,3} gather(f32[3,48,128]{2,1,0} %param_0.159, s32[3,4096,443,2]{3,2,1,0} %concatenate.37), 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(interpolate_backward)/jit(main)/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.72 = f32[3,48,4096,443]{2,1,0,3} multiply(f32[3,48,4096,443]{2,1,0,3} %broadcast.279, f32[3,48,4096,443]{2,1,0,3} %gather.37), metadata={op_name="jit(interpolate_backward)/jit(main)/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.275 = f32[3,48,4096,443]{2,1,0,3} broadcast(f32[3,4096,443]{2,1,0} %subtract.43), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %iota.34 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_157 = s32[] constant(1)
  %broadcast.273 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_157), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.91 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.23, s32[3,4096,443]{2,1,0} %broadcast.273), metadata={op_name="jit(interpolate_backward)/jit(main)/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.22 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.278, s32[3,4096,443]{2,1,0} %add.91, s32[3,4096,443]{2,1,0} %broadcast.277), metadata={op_name="jit(interpolate_backward)/jit(main)/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.38 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.22, s32[3,4096,443]{2,1,0} %broadcast.278), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.90 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.22, s32[3,4096,443]{2,1,0} %broadcast.276), metadata={op_name="jit(interpolate_backward)/jit(main)/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.38 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.38, s32[3,4096,443]{2,1,0} %add.90, s32[3,4096,443]{2,1,0} %clamp.22), metadata={op_name="jit(interpolate_backward)/jit(main)/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.62 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.38), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.36 = s32[3,4096,443,2]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.34, s32[3,4096,443,1]{2,1,0,3} %bitcast.62), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.36 = f32[3,48,4096,443]{2,1,0,3} gather(f32[3,48,128]{2,1,0} %param_0.159, s32[3,4096,443,2]{2,1,0,3} %concatenate.36), 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(interpolate_backward)/jit(main)/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.71 = f32[3,48,4096,443]{2,1,0,3} multiply(f32[3,48,4096,443]{2,1,0,3} %broadcast.275, f32[3,48,4096,443]{2,1,0,3} %gather.36), metadata={op_name="jit(interpolate_backward)/jit(main)/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.89 = f32[3,48,4096,443]{2,1,0,3} add(f32[3,48,4096,443]{2,1,0,3} %multiply.72, f32[3,48,4096,443]{2,1,0,3} %multiply.71), metadata={op_name="jit(interpolate_backward)/jit(main)/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.14 = f32[3,48,4096,443]{3,2,1,0} copy(f32[3,48,4096,443]{2,1,0,3} %add.89), metadata={op_name="jit(interpolate_backward)/jit(main)/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_4.128 = f32[3,4096,443]{2,1,0} parameter(4)
  %floor.119.clone.1 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_4.128), metadata={op_name="jit(interpolate_backward)/jit(main)/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.89.clone.1 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_4.128, f32[3,4096,443]{2,1,0} %floor.119.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.88.clone.1 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.280, f32[3,4096,443]{2,1,0} %subtract.89.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.196 = f32[3,4096,443]{2,1,0} parameter(3)
  %floor.118.clone.1 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_3.196), metadata={op_name="jit(interpolate_backward)/jit(main)/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.87.clone.1 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_3.196, f32[3,4096,443]{2,1,0} %floor.118.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.86.clone.1 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.280, f32[3,4096,443]{2,1,0} %subtract.87.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.91.clone.1 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.88.clone.1, f32[3,4096,443]{2,1,0} %subtract.86.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.482.clone.1 = f32[3,48,4096,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.91.clone.1), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.224 = f32[3,48,128,128]{3,2,1,0} parameter(2)
  %iota.54.clone.1 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.72.clone.1 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.119.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.71.clone.1 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.278, s32[3,4096,443]{2,1,0} %convert.72.clone.1, s32[3,4096,443]{2,1,0} %broadcast.277), metadata={op_name="jit(interpolate_backward)/jit(main)/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.89.clone.1 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.71.clone.1, s32[3,4096,443]{2,1,0} %broadcast.278), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.175.clone.1 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.71.clone.1, s32[3,4096,443]{2,1,0} %broadcast.276), metadata={op_name="jit(interpolate_backward)/jit(main)/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.89.clone.1 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.89.clone.1, s32[3,4096,443]{2,1,0} %add.175.clone.1, s32[3,4096,443]{2,1,0} %clamp.71.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.103.clone.1 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.89.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.71.clone.1 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.118.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.70.clone.1 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.278, s32[3,4096,443]{2,1,0} %convert.71.clone.1, s32[3,4096,443]{2,1,0} %broadcast.277), metadata={op_name="jit(interpolate_backward)/jit(main)/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.88.clone.1 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.70.clone.1, s32[3,4096,443]{2,1,0} %broadcast.278), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.174.clone.1 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.70.clone.1, s32[3,4096,443]{2,1,0} %broadcast.276), metadata={op_name="jit(interpolate_backward)/jit(main)/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.88.clone.1 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.88.clone.1, s32[3,4096,443]{2,1,0} %add.174.clone.1, s32[3,4096,443]{2,1,0} %clamp.70.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.102.clone.1 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.88.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.56.clone.1 = s32[3,4096,443,3]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.54.clone.1, s32[3,4096,443,1]{2,1,0,3} %bitcast.103.clone.1, s32[3,4096,443,1]{2,1,0,3} %bitcast.102.clone.1), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.17.clone.1 = f32[3,48,4096,443]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.224, s32[3,4096,443,3]{2,1,0,3} %concatenate.56.clone.1), 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(interpolate_backward)/jit(main)/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.51.clone.1 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %broadcast.482.clone.1, f32[3,48,4096,443]{3,2,1,0} %gather.17.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.87.clone.1 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.88.clone.1, f32[3,4096,443]{2,1,0} %subtract.87.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.434.clone.1 = f32[3,48,4096,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.87.clone.1), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.50.clone.1 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.95.clone.1 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.89.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.163.clone.1 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.71.clone.1, s32[3,4096,443]{2,1,0} %broadcast.273), metadata={op_name="jit(interpolate_backward)/jit(main)/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.62.clone.1 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.278, s32[3,4096,443]{2,1,0} %add.163.clone.1, s32[3,4096,443]{2,1,0} %broadcast.277), metadata={op_name="jit(interpolate_backward)/jit(main)/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.80.clone.1 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.62.clone.1, s32[3,4096,443]{2,1,0} %broadcast.278), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.162.clone.1 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.62.clone.1, s32[3,4096,443]{2,1,0} %broadcast.276), metadata={op_name="jit(interpolate_backward)/jit(main)/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.80.clone.1 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.80.clone.1, s32[3,4096,443]{2,1,0} %add.162.clone.1, s32[3,4096,443]{2,1,0} %clamp.62.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.94.clone.1 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.80.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.52.clone.1 = s32[3,4096,443,3]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.50.clone.1, s32[3,4096,443,1]{3,2,1,0} %bitcast.95.clone.1, s32[3,4096,443,1]{3,2,1,0} %bitcast.94.clone.1), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.16.clone.1 = f32[3,48,4096,443]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.224, s32[3,4096,443,3]{3,2,1,0} %concatenate.52.clone.1), 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(interpolate_backward)/jit(main)/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.50.clone.1 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %broadcast.434.clone.1, f32[3,48,4096,443]{3,2,1,0} %gather.16.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.53.clone.1 = f32[3,48,4096,443]{3,2,1,0} add(f32[3,48,4096,443]{3,2,1,0} %multiply.51.clone.1, f32[3,48,4096,443]{3,2,1,0} %multiply.50.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.83.clone.1 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.89.clone.1, f32[3,4096,443]{2,1,0} %subtract.86.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.386.clone.1 = f32[3,48,4096,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.83.clone.1), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.45.clone.1 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.146.clone.1 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.72.clone.1, s32[3,4096,443]{2,1,0} %broadcast.273), metadata={op_name="jit(interpolate_backward)/jit(main)/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.51.clone.1 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.278, s32[3,4096,443]{2,1,0} %add.146.clone.1, s32[3,4096,443]{2,1,0} %broadcast.277), metadata={op_name="jit(interpolate_backward)/jit(main)/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.69.clone.1 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.51.clone.1, s32[3,4096,443]{2,1,0} %broadcast.278), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.145.clone.1 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.51.clone.1, s32[3,4096,443]{2,1,0} %broadcast.276), metadata={op_name="jit(interpolate_backward)/jit(main)/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.69.clone.1 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.69.clone.1, s32[3,4096,443]{2,1,0} %add.145.clone.1, s32[3,4096,443]{2,1,0} %clamp.51.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.87.clone.1 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.69.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.47.clone.1 = s32[3,4096,443,3]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.45.clone.1, s32[3,4096,443,1]{2,1,0,3} %bitcast.87.clone.1, s32[3,4096,443,1]{2,1,0,3} %bitcast.102.clone.1), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.15.clone.1 = f32[3,48,4096,443]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.224, s32[3,4096,443,3]{2,1,0,3} %concatenate.47.clone.1), 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(interpolate_backward)/jit(main)/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.49.clone.1 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %broadcast.386.clone.1, f32[3,48,4096,443]{3,2,1,0} %gather.15.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.52.clone.1 = f32[3,48,4096,443]{3,2,1,0} add(f32[3,48,4096,443]{3,2,1,0} %add.53.clone.1, f32[3,48,4096,443]{3,2,1,0} %multiply.49.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.79.clone.1 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.89.clone.1, f32[3,4096,443]{2,1,0} %subtract.87.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.303.clone.1 = f32[3,48,4096,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.79.clone.1), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.41.clone.1 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.75.clone.1 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.69.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.43.clone.1 = s32[3,4096,443,3]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.41.clone.1, s32[3,4096,443,1]{3,2,1,0} %bitcast.75.clone.1, s32[3,4096,443,1]{3,2,1,0} %bitcast.94.clone.1), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.14.clone.1 = f32[3,48,4096,443]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.224, s32[3,4096,443,3]{3,2,1,0} %concatenate.43.clone.1), 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(interpolate_backward)/jit(main)/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.48.clone.1 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %broadcast.303.clone.1, f32[3,48,4096,443]{3,2,1,0} %gather.14.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.51.clone.1 = f32[3,48,4096,443]{3,2,1,0} add(f32[3,48,4096,443]{3,2,1,0} %add.52.clone.1, f32[3,48,4096,443]{3,2,1,0} %multiply.48.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.47 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %copy.14, f32[3,48,4096,443]{3,2,1,0} %add.51.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %exponential.1 = f32[3,48,4096,443]{3,2,1,0} exponential(f32[3,48,4096,443]{3,2,1,0} %multiply.47), metadata={op_name="jit(interpolate_backward)/jit(main)/exp" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=47}
  ROOT %tuple.6 = (f32[3,48,4096,443]{3,2,1,0}, f32[3,48,4096,443]{3,2,1,0}) tuple(f32[3,48,4096,443]{3,2,1,0} %exponential.1, f32[3,48,4096,443]{3,2,1,0} %add.51.clone.1)
}

%fused_computation.16 (param_0.103: f32[3,4096,443], param_1.132: s32[3,1], param_2.119: s32[3,1]) -> (f32[3,4096,443], f32[3,4096,443], f32[3,4096,443]) {
  %param_0.103 = f32[3,4096,443]{2,1,0} parameter(0)
  %bitcast.51 = f32[1,3,4096,443]{3,2,0,1} bitcast(f32[3,4096,443]{2,1,0} %param_0.103), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 443) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %param_2.119 = s32[3,1]{1,0} parameter(2)
  %gather.33 = f32[3,4096,443]{2,1,0} gather(f32[3,4096,443]{2,1,0} %param_0.103, s32[3,1]{1,0} %param_2.119), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,4096,443}, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443) 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.50 = f32[1,3,4096,443]{3,2,0,1} bitcast(f32[3,4096,443]{2,1,0} %gather.33), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 443) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %param_1.132 = s32[3,1]{1,0} parameter(1)
  %gather.32 = f32[3,4096,443]{2,1,0} gather(f32[3,4096,443]{2,1,0} %param_0.103, s32[3,1]{1,0} %param_1.132), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,4096,443}, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443) 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.49 = f32[1,3,4096,443]{3,2,0,1} bitcast(f32[3,4096,443]{2,1,0} %gather.32), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 443) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %concatenate.25 = f32[3,3,4096,443]{3,2,0,1} concatenate(f32[1,3,4096,443]{3,2,0,1} %bitcast.51, f32[1,3,4096,443]{3,2,0,1} %bitcast.50, f32[1,3,4096,443]{3,2,0,1} %bitcast.49), dimensions={0}, metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %constant_101 = f32[] constant(1)
  %broadcast.220 = f32[3,3,4096,443]{3,2,0,1} broadcast(f32[] %constant_101), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=159}
  %add.72 = f32[3,3,4096,443]{3,2,0,1} add(f32[3,3,4096,443]{3,2,0,1} %concatenate.25, f32[3,3,4096,443]{3,2,0,1} %broadcast.220), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=159}
  %constant_100 = f32[] constant(63.5)
  %broadcast.218 = f32[3,3,4096,443]{3,2,0,1} broadcast(f32[] %constant_100), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=162}
  %multiply.67 = f32[3,3,4096,443]{3,2,0,1} multiply(f32[3,3,4096,443]{3,2,0,1} %add.72, f32[3,3,4096,443]{3,2,0,1} %broadcast.218), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=162}
  %slice.9 = f32[3,1,4096,443]{3,2,0,1} slice(f32[3,3,4096,443]{3,2,0,1} %multiply.67), slice={[0:3], [1:2], [0:4096], [0:443]}, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.27 = f32[3,4096,443]{2,1,0} bitcast(f32[3,1,4096,443]{3,2,0,1} %slice.9), metadata={op_name="jit(interpolate_backward)/jit(main)/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.10.clone.1 = f32[3,1,4096,443]{3,2,0,1} slice(f32[3,3,4096,443]{3,2,0,1} %multiply.67), slice={[0:3], [0:1], [0:4096], [0:443]}, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443) 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.30.clone.1 = f32[3,4096,443]{2,1,0} bitcast(f32[3,1,4096,443]{3,2,0,1} %slice.10.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.8.clone.1 = f32[3,1,4096,443]{3,2,0,1} slice(f32[3,3,4096,443]{3,2,0,1} %multiply.67), slice={[0:3], [2:3], [0:4096], [0:443]}, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.26.clone.1 = f32[3,4096,443]{2,1,0} bitcast(f32[3,1,4096,443]{3,2,0,1} %slice.8.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.8 = (f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}) tuple(f32[3,4096,443]{2,1,0} %bitcast.27, f32[3,4096,443]{2,1,0} %bitcast.30.clone.1, f32[3,4096,443]{2,1,0} %bitcast.26.clone.1)
}

%fused_computation.23 (param_0.154: f32[3,48,128], param_1.217: f32[3,4096,443], param_2.218: f32[3,48,4096,443], param_3.184: f32[]) -> f32[3,48,4096,443] {
  %constant_179 = f32[] constant(1)
  %broadcast.301 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_179), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.217 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.51 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.217), metadata={op_name="jit(interpolate_backward)/jit(main)/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.49 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.217, f32[3,4096,443]{2,1,0} %floor.51), metadata={op_name="jit(interpolate_backward)/jit(main)/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.48 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.301, f32[3,4096,443]{2,1,0} %subtract.49), metadata={op_name="jit(interpolate_backward)/jit(main)/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.300 = f32[3,48,4096,443]{2,1,0,3} broadcast(f32[3,4096,443]{2,1,0} %subtract.48), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.154 = f32[3,48,128]{2,1,0} parameter(0)
  %iota.39 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_175 = s32[] constant(0)
  %broadcast.299 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_175), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.27 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.51), metadata={op_name="jit(interpolate_backward)/jit(main)/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_173 = s32[] constant(127)
  %broadcast.298 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_173), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.299, s32[3,4096,443]{2,1,0} %convert.27, s32[3,4096,443]{2,1,0} %broadcast.298), metadata={op_name="jit(interpolate_backward)/jit(main)/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.43 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.27, s32[3,4096,443]{2,1,0} %broadcast.299), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_172 = s32[] constant(128)
  %broadcast.297 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_172), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.100 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.27, s32[3,4096,443]{2,1,0} %broadcast.297), metadata={op_name="jit(interpolate_backward)/jit(main)/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.43 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.43, s32[3,4096,443]{2,1,0} %add.100, s32[3,4096,443]{2,1,0} %clamp.27), metadata={op_name="jit(interpolate_backward)/jit(main)/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.67 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.43), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.41 = s32[3,4096,443,2]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.39, s32[3,4096,443,1]{3,2,1,0} %bitcast.67), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.41 = f32[3,48,4096,443]{2,1,0,3} gather(f32[3,48,128]{2,1,0} %param_0.154, s32[3,4096,443,2]{3,2,1,0} %concatenate.41), 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(interpolate_backward)/jit(main)/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.77 = f32[3,48,4096,443]{2,1,0,3} multiply(f32[3,48,4096,443]{2,1,0,3} %broadcast.300, f32[3,48,4096,443]{2,1,0,3} %gather.41), metadata={op_name="jit(interpolate_backward)/jit(main)/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.296 = f32[3,48,4096,443]{2,1,0,3} broadcast(f32[3,4096,443]{2,1,0} %subtract.49), dimensions={0,2,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %iota.38 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_174 = s32[] constant(1)
  %broadcast.294 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_174), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.99 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.27, s32[3,4096,443]{2,1,0} %broadcast.294), metadata={op_name="jit(interpolate_backward)/jit(main)/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.26 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.299, s32[3,4096,443]{2,1,0} %add.99, s32[3,4096,443]{2,1,0} %broadcast.298), metadata={op_name="jit(interpolate_backward)/jit(main)/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.42 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.26, s32[3,4096,443]{2,1,0} %broadcast.299), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.98 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.26, s32[3,4096,443]{2,1,0} %broadcast.297), metadata={op_name="jit(interpolate_backward)/jit(main)/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.42 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.42, s32[3,4096,443]{2,1,0} %add.98, s32[3,4096,443]{2,1,0} %clamp.26), metadata={op_name="jit(interpolate_backward)/jit(main)/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.66 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.42), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.40 = s32[3,4096,443,2]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.38, s32[3,4096,443,1]{2,1,0,3} %bitcast.66), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.40 = f32[3,48,4096,443]{2,1,0,3} gather(f32[3,48,128]{2,1,0} %param_0.154, s32[3,4096,443,2]{2,1,0,3} %concatenate.40), 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(interpolate_backward)/jit(main)/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.76 = f32[3,48,4096,443]{2,1,0,3} multiply(f32[3,48,4096,443]{2,1,0,3} %broadcast.296, f32[3,48,4096,443]{2,1,0,3} %gather.40), metadata={op_name="jit(interpolate_backward)/jit(main)/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.97 = f32[3,48,4096,443]{2,1,0,3} add(f32[3,48,4096,443]{2,1,0,3} %multiply.77, f32[3,48,4096,443]{2,1,0,3} %multiply.76), metadata={op_name="jit(interpolate_backward)/jit(main)/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.16 = f32[3,48,4096,443]{3,2,1,0} copy(f32[3,48,4096,443]{2,1,0,3} %add.97), metadata={op_name="jit(interpolate_backward)/jit(main)/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_3.184 = f32[] parameter(3)
  %divide.3 = f32[] divide(f32[] %constant_179, f32[] %param_3.184), metadata={op_name="jit(interpolate_backward)/jit(main)/div" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=45}
  %broadcast.513 = f32[3,48,4096,443]{3,2,1,0} broadcast(f32[] %divide.3), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/broadcast_in_dim[shape=(144, 4096, 443) broadcast_dimensions=()]" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=46}
  %param_2.218 = f32[3,48,4096,443]{3,2,1,0} parameter(2)
  %multiply.95 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %broadcast.513, f32[3,48,4096,443]{3,2,1,0} %param_2.218), metadata={op_name="jit(interpolate_backward)/jit(main)/mul" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=47}
  ROOT %multiply.73 = f32[3,48,4096,443]{3,2,1,0} multiply(f32[3,48,4096,443]{3,2,1,0} %copy.16, f32[3,48,4096,443]{3,2,1,0} %multiply.95), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
}

ENTRY %main.395 (Arg_0.1: f32[3,48,128], Arg_1.2: f32[3,48,128,128], Arg_2.3: f32[3,4096,443]) -> (f32[3,48,128], f32[3,48,128,128]) {
  %Arg_0.1 = f32[3,48,128]{2,1,0} parameter(0)
  %Arg_2.3 = f32[3,4096,443]{2,1,0} parameter(2)
  %constant_51 = s32[3,1]{1,0} constant({ {1}, {2}, {0} }), metadata={op_name="jit(interpolate_backward)/jit(main)/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_47 = s32[3,1]{1,0} constant({ {2}, {0}, {1} }), metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %fusion.16 = (f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}) fusion(f32[3,4096,443]{2,1,0} %Arg_2.3, s32[3,1]{1,0} %constant_51, s32[3,1]{1,0} %constant_47), kind=kLoop, calls=%fused_computation.16, metadata={op_name="jit(interpolate_backward)/jit(main)/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.17 = f32[3,4096,443]{2,1,0} get-tuple-element((f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}) %fusion.16), index=1
  %Arg_1.2 = f32[3,48,128,128]{3,2,1,0} parameter(1)
  %get-tuple-element.18 = f32[3,4096,443]{2,1,0} get-tuple-element((f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}) %fusion.16), index=2
  %get-tuple-element.16 = f32[3,4096,443]{2,1,0} get-tuple-element((f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}) %fusion.16), index=0
  %fusion.1 = (f32[3,48,4096,443]{3,2,1,0}, f32[3,48,4096,443]{3,2,1,0}) fusion(f32[3,48,128]{2,1,0} %Arg_0.1, f32[3,4096,443]{2,1,0} %get-tuple-element.17, f32[3,48,128,128]{3,2,1,0} %Arg_1.2, f32[3,4096,443]{2,1,0} %get-tuple-element.18, f32[3,4096,443]{2,1,0} %get-tuple-element.16), kind=kLoop, calls=%fused_computation.1, metadata={op_name="jit(interpolate_backward)/jit(main)/exp" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=47}
  %get-tuple-element.15 = f32[3,48,4096,443]{3,2,1,0} get-tuple-element((f32[3,48,4096,443]{3,2,1,0}, f32[3,48,4096,443]{3,2,1,0}) %fusion.1), index=1
  %get-tuple-element.14 = f32[3,48,4096,443]{3,2,1,0} get-tuple-element((f32[3,48,4096,443]{3,2,1,0}, f32[3,48,4096,443]{3,2,1,0}) %fusion.1), index=0
  %input_fusion_reduce = f32[16165]{0} fusion(f32[3,48,4096,443]{3,2,1,0} %get-tuple-element.14), kind=kInput, calls=%input_fused_computation_reduce
  %constant_5 = f32[] constant(0)
  %reduce.287 = f32[] reduce(f32[16165]{0} %input_fusion_reduce, f32[] %constant_5), dimensions={0}, to_apply=%region_0.283, metadata={op_name="jit(interpolate_backward)/jit(main)/reduce_sum[axes=(0, 1, 2)]" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=46}
  %input_fusion_scatter.5 = f32[3,48,128]{2,1,0} fusion(f32[3,48,4096,443]{3,2,1,0} %get-tuple-element.15, f32[3,4096,443]{2,1,0} %get-tuple-element.17, f32[3,48,4096,443]{3,2,1,0} %get-tuple-element.14, f32[] %reduce.287), kind=kInput, calls=%input_fused_computation_scatter.5, metadata={op_name="jit(interpolate_backward)/jit(main)/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,48,128]{2,1,0} fusion(f32[3,48,4096,443]{3,2,1,0} %get-tuple-element.15, f32[3,4096,443]{2,1,0} %get-tuple-element.17, f32[3,48,4096,443]{3,2,1,0} %get-tuple-element.14, f32[] %reduce.287), kind=kInput, calls=%input_fused_computation_scatter.4, metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %add.41 = f32[3,48,128]{2,1,0} add(f32[3,48,128]{2,1,0} %input_fusion_scatter.5, f32[3,48,128]{2,1,0} %input_fusion_scatter.4), metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %fusion.23 = f32[3,48,4096,443]{3,2,1,0} fusion(f32[3,48,128]{2,1,0} %Arg_0.1, f32[3,4096,443]{2,1,0} %get-tuple-element.17, f32[3,48,4096,443]{3,2,1,0} %get-tuple-element.14, f32[] %reduce.287), kind=kLoop, calls=%fused_computation.23, metadata={op_name="jit(interpolate_backward)/jit(main)/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,48,128,128]{3,2,1,0} fusion(f32[3,48,4096,443]{3,2,1,0} %fusion.23, f32[3,4096,443]{2,1,0} %get-tuple-element.18, f32[3,4096,443]{2,1,0} %get-tuple-element.16), kind=kInput, calls=%input_fused_computation_scatter.3, metadata={op_name="jit(interpolate_backward)/jit(main)/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,48,128,128]{3,2,1,0} fusion(f32[3,48,4096,443]{3,2,1,0} %fusion.23, f32[3,4096,443]{2,1,0} %get-tuple-element.18, f32[3,4096,443]{2,1,0} %get-tuple-element.16), kind=kInput, calls=%input_fused_computation_scatter.2, metadata={op_name="jit(interpolate_backward)/jit(main)/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,48,128,128]{3,2,1,0} fusion(f32[3,48,4096,443]{3,2,1,0} %fusion.23, f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.3, f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.2, f32[3,4096,443]{2,1,0} %get-tuple-element.18, f32[3,4096,443]{2,1,0} %get-tuple-element.16), kind=kInput, calls=%input_fused_computation_scatter.1, metadata={op_name="jit(interpolate_backward)/jit(main)/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,48,128,128]{3,2,1,0} fusion(f32[3,48,4096,443]{3,2,1,0} %fusion.23, f32[3,4096,443]{2,1,0} %get-tuple-element.18, f32[3,4096,443]{2,1,0} %get-tuple-element.16), kind=kInput, calls=%input_fused_computation_scatter, metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %add.45 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.1, f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter), metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  ROOT %tuple.394 = (f32[3,48,128]{2,1,0}, f32[3,48,128,128]{3,2,1,0}) tuple(f32[3,48,128]{2,1,0} %add.41, f32[3,48,128,128]{3,2,1,0} %add.45)
}

更改后文本

打开文件

HloModule jit_interpolate_backward.17

%region_0.286 (Arg_0.287: f32[], Arg_1.288: f32[]) -> f32[] {
  %Arg_0.287 = f32[] parameter(0)
  %Arg_1.288 = f32[] parameter(1)
  ROOT %add.289 = f32[] add(f32[] %Arg_0.287, f32[] %Arg_1.288), metadata={op_name="jit(interpolate_backward)/jit(main)/reduce_sum[axes=(0, 1, 2)]" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=46}
}

%region_4.306 (Arg_0.307: f32[], Arg_1.308: f32[]) -> f32[] {
  %Arg_0.307 = f32[] parameter(0)
  %Arg_1.308 = f32[] parameter(1)
  ROOT %add.309 = f32[] add(f32[] %Arg_0.307, f32[] %Arg_1.308), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%input_fused_computation_scatter (param_0.154: f32[3,4096,48,443], param_1.215: f32[3,4096,443], param_2.216: f32[3,4096,443]) -> f32[3,48,128,128] {
  %constant_54 = f32[] constant(0)
  %broadcast.171 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_54), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.56 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_387 = s32[] constant(0)
  %broadcast.512 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_387), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.216 = f32[3,4096,443]{2,1,0} parameter(2)
  %floor.133 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_2.216), metadata={op_name="jit(interpolate_backward)/jit(main)/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.76 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.133), metadata={op_name="jit(interpolate_backward)/jit(main)/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_386 = s32[] constant(127)
  %broadcast.511 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_386), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.75 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.512, s32[3,4096,443]{2,1,0} %convert.76, s32[3,4096,443]{2,1,0} %broadcast.511), metadata={op_name="jit(interpolate_backward)/jit(main)/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.93 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.75, s32[3,4096,443]{2,1,0} %broadcast.512), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_385 = s32[] constant(128)
  %broadcast.510 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_385), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.180 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.75, s32[3,4096,443]{2,1,0} %broadcast.510), metadata={op_name="jit(interpolate_backward)/jit(main)/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.93 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.93, s32[3,4096,443]{2,1,0} %add.180, s32[3,4096,443]{2,1,0} %clamp.75), metadata={op_name="jit(interpolate_backward)/jit(main)/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.107 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.93), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.215 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.132 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.215), metadata={op_name="jit(interpolate_backward)/jit(main)/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.75 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.132), metadata={op_name="jit(interpolate_backward)/jit(main)/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.74 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.512, s32[3,4096,443]{2,1,0} %convert.75, s32[3,4096,443]{2,1,0} %broadcast.511), metadata={op_name="jit(interpolate_backward)/jit(main)/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.92 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.74, s32[3,4096,443]{2,1,0} %broadcast.512), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.179 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.74, s32[3,4096,443]{2,1,0} %broadcast.510), metadata={op_name="jit(interpolate_backward)/jit(main)/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.92 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.92, s32[3,4096,443]{2,1,0} %add.179, s32[3,4096,443]{2,1,0} %clamp.74), metadata={op_name="jit(interpolate_backward)/jit(main)/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.106 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.92), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.58 = s32[3,4096,443,3]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.56, s32[3,4096,443,1]{2,1,0,3} %bitcast.107, s32[3,4096,443,1]{2,1,0,3} %bitcast.106), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.8 = s32[3,4096,443,3]{3,2,1,0} copy(s32[3,4096,443,3]{2,1,0,3} %concatenate.58), metadata={op_name="jit(interpolate_backward)/jit(main)/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_364 = f32[] constant(1)
  %broadcast.488 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_364), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.97 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_2.216, f32[3,4096,443]{2,1,0} %floor.133), metadata={op_name="jit(interpolate_backward)/jit(main)/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.96 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.488, f32[3,4096,443]{2,1,0} %subtract.97), metadata={op_name="jit(interpolate_backward)/jit(main)/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.95 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.215, f32[3,4096,443]{2,1,0} %floor.132), metadata={op_name="jit(interpolate_backward)/jit(main)/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.94 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.488, f32[3,4096,443]{2,1,0} %subtract.95), metadata={op_name="jit(interpolate_backward)/jit(main)/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.94 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.96, f32[3,4096,443]{2,1,0} %subtract.94), metadata={op_name="jit(interpolate_backward)/jit(main)/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.487 = f32[3,4096,48,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.94), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.154 = f32[3,4096,48,443]{3,2,1,0} parameter(0)
  %multiply.41 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %broadcast.487, f32[3,4096,48,443]{3,2,1,0} %param_0.154), metadata={op_name="jit(interpolate_backward)/jit(main)/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.12 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.171, s32[3,4096,443,3]{3,2,1,0} %copy.8, f32[3,4096,48,443]{3,2,1,0} %multiply.41), 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_4.306, metadata={op_name="jit(interpolate_backward)/jit(main)/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_3.302 (Arg_0.303: f32[], Arg_1.304: f32[]) -> f32[] {
  %Arg_0.303 = f32[] parameter(0)
  %Arg_1.304 = f32[] parameter(1)
  ROOT %add.305 = f32[] add(f32[] %Arg_0.303, f32[] %Arg_1.304), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%input_fused_computation_scatter.1 (param_0.151: f32[3,4096,48,443], param_1.207: f32[3,48,128,128], param_2.209: f32[3,48,128,128], param_3.173: f32[3,4096,443], param_4.100: f32[3,4096,443]) -> f32[3,48,128,128] {
  %param_1.207 = f32[3,48,128,128]{3,2,1,0} parameter(1)
  %param_2.209 = f32[3,48,128,128]{3,2,1,0} parameter(2)
  %add.50 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %param_1.207, f32[3,48,128,128]{3,2,1,0} %param_2.209), metadata={op_name="jit(interpolate_backward)/jit(main)/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.52 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_354 = s32[] constant(0)
  %broadcast.480 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_354), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.100 = f32[3,4096,443]{2,1,0} parameter(4)
  %floor.115 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_4.100), metadata={op_name="jit(interpolate_backward)/jit(main)/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.68 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.115), metadata={op_name="jit(interpolate_backward)/jit(main)/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_352 = s32[] constant(127)
  %broadcast.479 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_352), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.480, s32[3,4096,443]{2,1,0} %convert.68, s32[3,4096,443]{2,1,0} %broadcast.479), metadata={op_name="jit(interpolate_backward)/jit(main)/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.85 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.67, s32[3,4096,443]{2,1,0} %broadcast.480), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_351 = s32[] constant(128)
  %broadcast.478 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_351), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.170 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.67, s32[3,4096,443]{2,1,0} %broadcast.478), metadata={op_name="jit(interpolate_backward)/jit(main)/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.85 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.85, s32[3,4096,443]{2,1,0} %add.170, s32[3,4096,443]{2,1,0} %clamp.67), metadata={op_name="jit(interpolate_backward)/jit(main)/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.99 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.85), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_3.173 = f32[3,4096,443]{2,1,0} parameter(3)
  %floor.114 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_3.173), metadata={op_name="jit(interpolate_backward)/jit(main)/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.67 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.114), metadata={op_name="jit(interpolate_backward)/jit(main)/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_353 = s32[] constant(1)
  %broadcast.476 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_353), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.169 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.67, s32[3,4096,443]{2,1,0} %broadcast.476), metadata={op_name="jit(interpolate_backward)/jit(main)/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.66 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.480, s32[3,4096,443]{2,1,0} %add.169, s32[3,4096,443]{2,1,0} %broadcast.479), metadata={op_name="jit(interpolate_backward)/jit(main)/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.84 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.66, s32[3,4096,443]{2,1,0} %broadcast.480), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.168 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.66, s32[3,4096,443]{2,1,0} %broadcast.478), metadata={op_name="jit(interpolate_backward)/jit(main)/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.84 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.84, s32[3,4096,443]{2,1,0} %add.168, s32[3,4096,443]{2,1,0} %clamp.66), metadata={op_name="jit(interpolate_backward)/jit(main)/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.98 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.84), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.54 = s32[3,4096,443,3]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.52, s32[3,4096,443,1]{3,2,1,0} %bitcast.99, s32[3,4096,443,1]{3,2,1,0} %bitcast.98), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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_312 = f32[] constant(1)
  %broadcast.440 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_312), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.81 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_4.100, f32[3,4096,443]{2,1,0} %floor.115), metadata={op_name="jit(interpolate_backward)/jit(main)/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.80 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.440, f32[3,4096,443]{2,1,0} %subtract.81), metadata={op_name="jit(interpolate_backward)/jit(main)/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.79 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_3.173, f32[3,4096,443]{2,1,0} %floor.114), metadata={op_name="jit(interpolate_backward)/jit(main)/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.90 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.80, f32[3,4096,443]{2,1,0} %subtract.79), metadata={op_name="jit(interpolate_backward)/jit(main)/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.439 = f32[3,4096,48,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.90), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.151 = f32[3,4096,48,443]{3,2,1,0} parameter(0)
  %multiply.42 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %broadcast.439, f32[3,4096,48,443]{3,2,1,0} %param_0.151), metadata={op_name="jit(interpolate_backward)/jit(main)/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.13 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %add.50, s32[3,4096,443,3]{3,2,1,0} %concatenate.54, f32[3,4096,48,443]{3,2,1,0} %multiply.42), 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_3.302, metadata={op_name="jit(interpolate_backward)/jit(main)/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_2.298 (Arg_0.299: f32[], Arg_1.300: f32[]) -> f32[] {
  %Arg_0.299 = f32[] parameter(0)
  %Arg_1.300 = f32[] parameter(1)
  ROOT %add.301 = f32[] add(f32[] %Arg_0.299, f32[] %Arg_1.300), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%input_fused_computation_scatter.2 (param_0.136: f32[3,4096,48,443], param_1.191: f32[3,4096,443], param_2.191: f32[3,4096,443]) -> f32[3,48,128,128] {
  %constant_55 = f32[] constant(0)
  %broadcast.172 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_55), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.47 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_302 = s32[] constant(0)
  %broadcast.432 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_302), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.191 = f32[3,4096,443]{2,1,0} parameter(2)
  %floor.95 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_2.191), metadata={op_name="jit(interpolate_backward)/jit(main)/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.56 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.95), metadata={op_name="jit(interpolate_backward)/jit(main)/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_305 = s32[] constant(1)
  %broadcast.431 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_305), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.152 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.56, s32[3,4096,443]{2,1,0} %broadcast.431), metadata={op_name="jit(interpolate_backward)/jit(main)/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_301 = s32[] constant(127)
  %broadcast.430 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_301), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.432, s32[3,4096,443]{2,1,0} %add.152, s32[3,4096,443]{2,1,0} %broadcast.430), metadata={op_name="jit(interpolate_backward)/jit(main)/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.73 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.55, s32[3,4096,443]{2,1,0} %broadcast.432), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_300 = s32[] constant(128)
  %broadcast.429 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_300), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.151 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.55, s32[3,4096,443]{2,1,0} %broadcast.429), metadata={op_name="jit(interpolate_backward)/jit(main)/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.73 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.73, s32[3,4096,443]{2,1,0} %add.151, s32[3,4096,443]{2,1,0} %clamp.55), metadata={op_name="jit(interpolate_backward)/jit(main)/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.91 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.73), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.191 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.94 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.191), metadata={op_name="jit(interpolate_backward)/jit(main)/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.55 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.94), metadata={op_name="jit(interpolate_backward)/jit(main)/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.54 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.432, s32[3,4096,443]{2,1,0} %convert.55, s32[3,4096,443]{2,1,0} %broadcast.430), metadata={op_name="jit(interpolate_backward)/jit(main)/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.72 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.54, s32[3,4096,443]{2,1,0} %broadcast.432), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.150 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.54, s32[3,4096,443]{2,1,0} %broadcast.429), metadata={op_name="jit(interpolate_backward)/jit(main)/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.72 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.72, s32[3,4096,443]{2,1,0} %add.150, s32[3,4096,443]{2,1,0} %clamp.54), metadata={op_name="jit(interpolate_backward)/jit(main)/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.90 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.72), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.49 = s32[3,4096,443,3]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.47, s32[3,4096,443,1]{2,1,0,3} %bitcast.91, s32[3,4096,443,1]{2,1,0,3} %bitcast.90), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.9 = s32[3,4096,443,3]{3,2,1,0} copy(s32[3,4096,443,3]{2,1,0,3} %concatenate.49), metadata={op_name="jit(interpolate_backward)/jit(main)/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.69 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_2.191, f32[3,4096,443]{2,1,0} %floor.95), metadata={op_name="jit(interpolate_backward)/jit(main)/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_262 = f32[] constant(1)
  %broadcast.392 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_262), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.68 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.191, f32[3,4096,443]{2,1,0} %floor.94), metadata={op_name="jit(interpolate_backward)/jit(main)/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.67 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.392, f32[3,4096,443]{2,1,0} %subtract.68), metadata={op_name="jit(interpolate_backward)/jit(main)/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.86 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.69, f32[3,4096,443]{2,1,0} %subtract.67), metadata={op_name="jit(interpolate_backward)/jit(main)/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.391 = f32[3,4096,48,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.86), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.136 = f32[3,4096,48,443]{3,2,1,0} parameter(0)
  %multiply.43 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %broadcast.391, f32[3,4096,48,443]{3,2,1,0} %param_0.136), metadata={op_name="jit(interpolate_backward)/jit(main)/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.14 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.172, s32[3,4096,443,3]{3,2,1,0} %copy.9, f32[3,4096,48,443]{3,2,1,0} %multiply.43), 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_2.298, metadata={op_name="jit(interpolate_backward)/jit(main)/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_1.294 (Arg_0.295: f32[], Arg_1.296: f32[]) -> f32[] {
  %Arg_0.295 = f32[] parameter(0)
  %Arg_1.296 = f32[] parameter(1)
  ROOT %add.297 = f32[] add(f32[] %Arg_0.295, f32[] %Arg_1.296), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%input_fused_computation_scatter.3 (param_0.133: f32[3,4096,48,443], param_1.173: f32[3,4096,443], param_2.171: f32[3,4096,443]) -> f32[3,48,128,128] {
  %constant_56 = f32[] constant(0)
  %broadcast.173 = f32[3,48,128,128]{3,2,1,0} broadcast(f32[] %constant_56), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.43 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_252 = s32[] constant(0)
  %broadcast.384 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_252), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.171 = f32[3,4096,443]{2,1,0} parameter(2)
  %floor.75 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_2.171), metadata={op_name="jit(interpolate_backward)/jit(main)/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.44 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.75), metadata={op_name="jit(interpolate_backward)/jit(main)/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_251 = s32[] constant(1)
  %broadcast.383 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_251), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.134 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.44, s32[3,4096,443]{2,1,0} %broadcast.383), metadata={op_name="jit(interpolate_backward)/jit(main)/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_250 = s32[] constant(127)
  %broadcast.382 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_250), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.43 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.384, s32[3,4096,443]{2,1,0} %add.134, s32[3,4096,443]{2,1,0} %broadcast.382), metadata={op_name="jit(interpolate_backward)/jit(main)/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.61 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.43, s32[3,4096,443]{2,1,0} %broadcast.384), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_249 = s32[] constant(128)
  %broadcast.381 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_249), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.132 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.43, s32[3,4096,443]{2,1,0} %broadcast.381), metadata={op_name="jit(interpolate_backward)/jit(main)/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.61 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.61, s32[3,4096,443]{2,1,0} %add.132, s32[3,4096,443]{2,1,0} %clamp.43), metadata={op_name="jit(interpolate_backward)/jit(main)/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.79 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.61), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %param_1.173 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.74 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.173), metadata={op_name="jit(interpolate_backward)/jit(main)/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.43 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.74), metadata={op_name="jit(interpolate_backward)/jit(main)/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.131 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.43, s32[3,4096,443]{2,1,0} %broadcast.383), metadata={op_name="jit(interpolate_backward)/jit(main)/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.42 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.384, s32[3,4096,443]{2,1,0} %add.131, s32[3,4096,443]{2,1,0} %broadcast.382), metadata={op_name="jit(interpolate_backward)/jit(main)/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.60 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.42, s32[3,4096,443]{2,1,0} %broadcast.384), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.130 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.42, s32[3,4096,443]{2,1,0} %broadcast.381), metadata={op_name="jit(interpolate_backward)/jit(main)/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.59 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.60, s32[3,4096,443]{2,1,0} %add.130, s32[3,4096,443]{2,1,0} %clamp.42), metadata={op_name="jit(interpolate_backward)/jit(main)/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.78 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.59), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.45 = s32[3,4096,443,3]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.43, s32[3,4096,443,1]{3,2,1,0} %bitcast.79, s32[3,4096,443,1]{3,2,1,0} %bitcast.78), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.57 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_2.171, f32[3,4096,443]{2,1,0} %floor.75), metadata={op_name="jit(interpolate_backward)/jit(main)/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.56 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.173, f32[3,4096,443]{2,1,0} %floor.74), metadata={op_name="jit(interpolate_backward)/jit(main)/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.82 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.57, f32[3,4096,443]{2,1,0} %subtract.56), metadata={op_name="jit(interpolate_backward)/jit(main)/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.306 = f32[3,4096,48,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.82), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.133 = f32[3,4096,48,443]{3,2,1,0} parameter(0)
  %multiply.44 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %broadcast.306, f32[3,4096,48,443]{3,2,1,0} %param_0.133), metadata={op_name="jit(interpolate_backward)/jit(main)/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.15 = f32[3,48,128,128]{3,2,1,0} scatter(f32[3,48,128,128]{3,2,1,0} %broadcast.173, s32[3,4096,443,3]{3,2,1,0} %concatenate.45, f32[3,4096,48,443]{3,2,1,0} %multiply.44), 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_1.294, metadata={op_name="jit(interpolate_backward)/jit(main)/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_6.349 (Arg_0.350: f32[], Arg_1.351: f32[]) -> f32[] {
  %Arg_0.350 = f32[] parameter(0)
  %Arg_1.351 = f32[] parameter(1)
  ROOT %add.352 = f32[] add(f32[] %Arg_0.350, f32[] %Arg_1.351), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%input_fused_computation_scatter.4 (param_0.162: f32[3,4096,48,443], param_1.221: f32[3,4096,443], param_2.222: f32[3,4096,48,443], param_3.192: f32[]) -> f32[3,48,128] {
  %constant_57 = f32[] constant(0)
  %broadcast.174 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_57), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.31 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_145 = s32[] constant(0)
  %broadcast.261 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_145), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.221 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.34 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.221), metadata={op_name="jit(interpolate_backward)/jit(main)/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.19 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.34), metadata={op_name="jit(interpolate_backward)/jit(main)/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_144 = s32[] constant(127)
  %broadcast.260 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_144), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.19 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.261, s32[3,4096,443]{2,1,0} %convert.19, s32[3,4096,443]{2,1,0} %broadcast.260), metadata={op_name="jit(interpolate_backward)/jit(main)/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.35 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.19, s32[3,4096,443]{2,1,0} %broadcast.261), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_143 = s32[] constant(128)
  %broadcast.259 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_143), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.84 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.19, s32[3,4096,443]{2,1,0} %broadcast.259), metadata={op_name="jit(interpolate_backward)/jit(main)/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.35 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.35, s32[3,4096,443]{2,1,0} %add.84, s32[3,4096,443]{2,1,0} %clamp.19), metadata={op_name="jit(interpolate_backward)/jit(main)/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.59 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.35), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.33 = s32[3,4096,443,2]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.31, s32[3,4096,443,1]{3,2,1,0} %bitcast.59), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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_131 = f32[] constant(1)
  %broadcast.249 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_131), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.35 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.221, f32[3,4096,443]{2,1,0} %floor.34), metadata={op_name="jit(interpolate_backward)/jit(main)/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.34 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.249, f32[3,4096,443]{2,1,0} %subtract.35), metadata={op_name="jit(interpolate_backward)/jit(main)/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.248 = f32[3,4096,48,443]{2,1,0,3} broadcast(f32[3,4096,443]{2,1,0} %subtract.34), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.10 = f32[3,4096,48,443]{3,2,1,0} copy(f32[3,4096,48,443]{2,1,0,3} %broadcast.248), metadata={op_name="jit(interpolate_backward)/jit(main)/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_3.192 = f32[] parameter(3)
  %divide.7 = f32[] divide(f32[] %constant_131, f32[] %param_3.192), metadata={op_name="jit(interpolate_backward)/jit(main)/div" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=45}
  %broadcast.518 = f32[3,4096,48,443]{3,2,1,0} broadcast(f32[] %divide.7), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/broadcast_in_dim[shape=(144, 4096, 443) broadcast_dimensions=()]" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=46}
  %param_2.222 = f32[3,4096,48,443]{3,2,1,0} parameter(2)
  %multiply.100 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %broadcast.518, f32[3,4096,48,443]{3,2,1,0} %param_2.222), metadata={op_name="jit(interpolate_backward)/jit(main)/mul" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=47}
  %param_0.162 = f32[3,4096,48,443]{3,2,1,0} parameter(0)
  %multiply.61 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %multiply.100, f32[3,4096,48,443]{3,2,1,0} %param_0.162), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %multiply.45 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %copy.10, f32[3,4096,48,443]{3,2,1,0} %multiply.61), metadata={op_name="jit(interpolate_backward)/jit(main)/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.16 = f32[3,48,128]{2,1,0} scatter(f32[3,48,128]{2,1,0} %broadcast.174, s32[3,4096,443,2]{3,2,1,0} %concatenate.33, f32[3,4096,48,443]{3,2,1,0} %multiply.45), update_window_dims={2}, inserted_window_dims={0,2}, scatter_dims_to_operand_dims={0,2}, index_vector_dim=3, to_apply=%region_6.349, metadata={op_name="jit(interpolate_backward)/jit(main)/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_5.345 (Arg_0.346: f32[], Arg_1.347: f32[]) -> f32[] {
  %Arg_0.346 = f32[] parameter(0)
  %Arg_1.347 = f32[] parameter(1)
  ROOT %add.348 = f32[] add(f32[] %Arg_0.346, f32[] %Arg_1.347), metadata={op_name="add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
}

%input_fused_computation_scatter.5 (param_0.160: f32[3,4096,48,443], param_1.219: f32[3,4096,443], param_2.220: f32[3,4096,48,443], param_3.188: f32[]) -> f32[3,48,128] {
  %constant_58 = f32[] constant(0)
  %broadcast.176 = f32[3,48,128]{2,1,0} broadcast(f32[] %constant_58), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.27 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_125 = s32[] constant(0)
  %broadcast.241 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_125), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.219 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.26 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.219), metadata={op_name="jit(interpolate_backward)/jit(main)/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.15 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.26), metadata={op_name="jit(interpolate_backward)/jit(main)/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_124 = s32[] constant(1)
  %broadcast.240 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_124), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.80 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.15, s32[3,4096,443]{2,1,0} %broadcast.240), metadata={op_name="jit(interpolate_backward)/jit(main)/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_123 = s32[] constant(127)
  %broadcast.239 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_123), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.15 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.241, s32[3,4096,443]{2,1,0} %add.80, s32[3,4096,443]{2,1,0} %broadcast.239), metadata={op_name="jit(interpolate_backward)/jit(main)/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.31 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.15, s32[3,4096,443]{2,1,0} %broadcast.241), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_122 = s32[] constant(128)
  %broadcast.238 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_122), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.79 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.15, s32[3,4096,443]{2,1,0} %broadcast.238), metadata={op_name="jit(interpolate_backward)/jit(main)/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.31 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.31, s32[3,4096,443]{2,1,0} %add.79, s32[3,4096,443]{2,1,0} %clamp.15), metadata={op_name="jit(interpolate_backward)/jit(main)/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.55 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.31), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.29 = s32[3,4096,443,2]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.27, s32[3,4096,443,1]{2,1,0,3} %bitcast.55), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.13 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.219, f32[3,4096,443]{2,1,0} %floor.26), metadata={op_name="jit(interpolate_backward)/jit(main)/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.175 = f32[3,4096,48,443]{2,1,0,3} broadcast(f32[3,4096,443]{2,1,0} %subtract.13), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %constant_396 = f32[] constant(1)
  %param_3.188 = f32[] parameter(3)
  %divide.5 = f32[] divide(f32[] %constant_396, f32[] %param_3.188), metadata={op_name="jit(interpolate_backward)/jit(main)/div" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=45}
  %broadcast.516 = f32[3,4096,48,443]{3,2,1,0} broadcast(f32[] %divide.5), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/broadcast_in_dim[shape=(144, 4096, 443) broadcast_dimensions=()]" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=46}
  %param_2.220 = f32[3,4096,48,443]{3,2,1,0} parameter(2)
  %multiply.98 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %broadcast.516, f32[3,4096,48,443]{3,2,1,0} %param_2.220), metadata={op_name="jit(interpolate_backward)/jit(main)/mul" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=47}
  %param_0.160 = f32[3,4096,48,443]{3,2,1,0} parameter(0)
  %multiply.62 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %multiply.98, f32[3,4096,48,443]{3,2,1,0} %param_0.160), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %copy.11 = f32[3,4096,48,443]{2,1,0,3} copy(f32[3,4096,48,443]{3,2,1,0} %multiply.62), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %multiply.46 = f32[3,4096,48,443]{2,1,0,3} multiply(f32[3,4096,48,443]{2,1,0,3} %broadcast.175, f32[3,4096,48,443]{2,1,0,3} %copy.11), metadata={op_name="jit(interpolate_backward)/jit(main)/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.17 = f32[3,48,128]{2,1,0} scatter(f32[3,48,128]{2,1,0} %broadcast.176, s32[3,4096,443,2]{2,1,0,3} %concatenate.29, f32[3,4096,48,443]{2,1,0,3} %multiply.46), update_window_dims={2}, inserted_window_dims={0,2}, scatter_dims_to_operand_dims={0,2}, index_vector_dim=3, to_apply=%region_5.345, metadata={op_name="jit(interpolate_backward)/jit(main)/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_fused_computation_reduce (param_0.69: f32[3,4096,48,443]) -> f32[16165] {
  %param_0.69 = f32[3,4096,48,443]{3,2,1,0} parameter(0)
  %bitcast.19 = f32[261292032]{0} bitcast(f32[3,4096,48,443]{3,2,1,0} %param_0.69)
  %constant_60 = f32[] constant(0)
  %pad.1 = f32[261307225]{0} pad(f32[261292032]{0} %bitcast.19, f32[] %constant_60), padding=0_15193
  %bitcast.18 = f32[16165,16165]{1,0} bitcast(f32[261307225]{0} %pad.1)
  ROOT %reduce.5 = f32[16165]{0} reduce(f32[16165,16165]{1,0} %bitcast.18, f32[] %constant_60), dimensions={1}, to_apply=%region_0.286
}

%fused_computation.1 (param_0.163: f32[3,48,128], param_1.222: f32[3,4096,443], param_2.224: f32[3,48,128,128], param_3.196: f32[3,4096,443], param_4.128: f32[3,4096,443]) -> (f32[3,4096,48,443], f32[3,4096,48,443]) {
  %constant_162 = f32[] constant(1)
  %broadcast.281 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_162), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.222 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.43 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.222), metadata={op_name="jit(interpolate_backward)/jit(main)/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.43 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.222, f32[3,4096,443]{2,1,0} %floor.43), metadata={op_name="jit(interpolate_backward)/jit(main)/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.42 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.281, f32[3,4096,443]{2,1,0} %subtract.43), metadata={op_name="jit(interpolate_backward)/jit(main)/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.280 = f32[3,4096,48,443]{2,1,0,3} broadcast(f32[3,4096,443]{2,1,0} %subtract.42), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.163 = f32[3,48,128]{2,1,0} parameter(0)
  %iota.35 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_158 = s32[] constant(0)
  %broadcast.279 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_158), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.23 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.43), metadata={op_name="jit(interpolate_backward)/jit(main)/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_156 = s32[] constant(127)
  %broadcast.278 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_156), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.23 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.279, s32[3,4096,443]{2,1,0} %convert.23, s32[3,4096,443]{2,1,0} %broadcast.278), metadata={op_name="jit(interpolate_backward)/jit(main)/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.39 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.23, s32[3,4096,443]{2,1,0} %broadcast.279), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_155 = s32[] constant(128)
  %broadcast.277 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_155), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.92 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.23, s32[3,4096,443]{2,1,0} %broadcast.277), metadata={op_name="jit(interpolate_backward)/jit(main)/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.39 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.39, s32[3,4096,443]{2,1,0} %add.92, s32[3,4096,443]{2,1,0} %clamp.23), metadata={op_name="jit(interpolate_backward)/jit(main)/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.63 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.39), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.37 = s32[3,4096,443,2]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.35, s32[3,4096,443,1]{3,2,1,0} %bitcast.63), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.37 = f32[3,4096,48,443]{2,1,0,3} gather(f32[3,48,128]{2,1,0} %param_0.163, s32[3,4096,443,2]{3,2,1,0} %concatenate.37), 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(interpolate_backward)/jit(main)/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.73 = f32[3,4096,48,443]{2,1,0,3} multiply(f32[3,4096,48,443]{2,1,0,3} %broadcast.280, f32[3,4096,48,443]{2,1,0,3} %gather.37), metadata={op_name="jit(interpolate_backward)/jit(main)/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.276 = f32[3,4096,48,443]{2,1,0,3} broadcast(f32[3,4096,443]{2,1,0} %subtract.43), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %iota.34 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_157 = s32[] constant(1)
  %broadcast.274 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_157), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.91 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.23, s32[3,4096,443]{2,1,0} %broadcast.274), metadata={op_name="jit(interpolate_backward)/jit(main)/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.22 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.279, s32[3,4096,443]{2,1,0} %add.91, s32[3,4096,443]{2,1,0} %broadcast.278), metadata={op_name="jit(interpolate_backward)/jit(main)/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.38 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.22, s32[3,4096,443]{2,1,0} %broadcast.279), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.90 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.22, s32[3,4096,443]{2,1,0} %broadcast.277), metadata={op_name="jit(interpolate_backward)/jit(main)/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.38 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.38, s32[3,4096,443]{2,1,0} %add.90, s32[3,4096,443]{2,1,0} %clamp.22), metadata={op_name="jit(interpolate_backward)/jit(main)/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.62 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.38), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.36 = s32[3,4096,443,2]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.34, s32[3,4096,443,1]{2,1,0,3} %bitcast.62), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.36 = f32[3,4096,48,443]{2,1,0,3} gather(f32[3,48,128]{2,1,0} %param_0.163, s32[3,4096,443,2]{2,1,0,3} %concatenate.36), 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(interpolate_backward)/jit(main)/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.72 = f32[3,4096,48,443]{2,1,0,3} multiply(f32[3,4096,48,443]{2,1,0,3} %broadcast.276, f32[3,4096,48,443]{2,1,0,3} %gather.36), metadata={op_name="jit(interpolate_backward)/jit(main)/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.89 = f32[3,4096,48,443]{2,1,0,3} add(f32[3,4096,48,443]{2,1,0,3} %multiply.73, f32[3,4096,48,443]{2,1,0,3} %multiply.72), metadata={op_name="jit(interpolate_backward)/jit(main)/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.14 = f32[3,4096,48,443]{3,2,1,0} copy(f32[3,4096,48,443]{2,1,0,3} %add.89), metadata={op_name="jit(interpolate_backward)/jit(main)/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.128 = f32[3,4096,443]{2,1,0} parameter(4)
  %floor.119.clone.1 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_4.128), metadata={op_name="jit(interpolate_backward)/jit(main)/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.89.clone.1 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_4.128, f32[3,4096,443]{2,1,0} %floor.119.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.88.clone.1 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.281, f32[3,4096,443]{2,1,0} %subtract.89.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.196 = f32[3,4096,443]{2,1,0} parameter(3)
  %floor.118.clone.1 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_3.196), metadata={op_name="jit(interpolate_backward)/jit(main)/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.87.clone.1 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_3.196, f32[3,4096,443]{2,1,0} %floor.118.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.86.clone.1 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.281, f32[3,4096,443]{2,1,0} %subtract.87.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.92.clone.1 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.88.clone.1, f32[3,4096,443]{2,1,0} %subtract.86.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.483.clone.1 = f32[3,4096,48,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.92.clone.1), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.224 = f32[3,48,128,128]{3,2,1,0} parameter(2)
  %iota.54.clone.1 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.72.clone.1 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.119.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.71.clone.1 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.279, s32[3,4096,443]{2,1,0} %convert.72.clone.1, s32[3,4096,443]{2,1,0} %broadcast.278), metadata={op_name="jit(interpolate_backward)/jit(main)/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.89.clone.1 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.71.clone.1, s32[3,4096,443]{2,1,0} %broadcast.279), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.175.clone.1 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.71.clone.1, s32[3,4096,443]{2,1,0} %broadcast.277), metadata={op_name="jit(interpolate_backward)/jit(main)/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.89.clone.1 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.89.clone.1, s32[3,4096,443]{2,1,0} %add.175.clone.1, s32[3,4096,443]{2,1,0} %clamp.71.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.103.clone.1 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.89.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %convert.71.clone.1 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.118.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.70.clone.1 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.279, s32[3,4096,443]{2,1,0} %convert.71.clone.1, s32[3,4096,443]{2,1,0} %broadcast.278), metadata={op_name="jit(interpolate_backward)/jit(main)/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.88.clone.1 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.70.clone.1, s32[3,4096,443]{2,1,0} %broadcast.279), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.174.clone.1 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.70.clone.1, s32[3,4096,443]{2,1,0} %broadcast.277), metadata={op_name="jit(interpolate_backward)/jit(main)/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.88.clone.1 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.88.clone.1, s32[3,4096,443]{2,1,0} %add.174.clone.1, s32[3,4096,443]{2,1,0} %clamp.70.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.102.clone.1 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.88.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.56.clone.1 = s32[3,4096,443,3]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.54.clone.1, s32[3,4096,443,1]{2,1,0,3} %bitcast.103.clone.1, s32[3,4096,443,1]{2,1,0,3} %bitcast.102.clone.1), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.17.clone.1 = f32[3,4096,48,443]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.224, s32[3,4096,443,3]{2,1,0,3} %concatenate.56.clone.1), 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(interpolate_backward)/jit(main)/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.52.clone.1 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %broadcast.483.clone.1, f32[3,4096,48,443]{3,2,1,0} %gather.17.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.88.clone.1 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.88.clone.1, f32[3,4096,443]{2,1,0} %subtract.87.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.435.clone.1 = f32[3,4096,48,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.88.clone.1), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.50.clone.1 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.95.clone.1 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.89.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.163.clone.1 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.71.clone.1, s32[3,4096,443]{2,1,0} %broadcast.274), metadata={op_name="jit(interpolate_backward)/jit(main)/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.62.clone.1 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.279, s32[3,4096,443]{2,1,0} %add.163.clone.1, s32[3,4096,443]{2,1,0} %broadcast.278), metadata={op_name="jit(interpolate_backward)/jit(main)/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.80.clone.1 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.62.clone.1, s32[3,4096,443]{2,1,0} %broadcast.279), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.162.clone.1 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.62.clone.1, s32[3,4096,443]{2,1,0} %broadcast.277), metadata={op_name="jit(interpolate_backward)/jit(main)/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.80.clone.1 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.80.clone.1, s32[3,4096,443]{2,1,0} %add.162.clone.1, s32[3,4096,443]{2,1,0} %clamp.62.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.94.clone.1 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.80.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.52.clone.1 = s32[3,4096,443,3]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.50.clone.1, s32[3,4096,443,1]{3,2,1,0} %bitcast.95.clone.1, s32[3,4096,443,1]{3,2,1,0} %bitcast.94.clone.1), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.16.clone.1 = f32[3,4096,48,443]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.224, s32[3,4096,443,3]{3,2,1,0} %concatenate.52.clone.1), 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(interpolate_backward)/jit(main)/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.51.clone.1 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %broadcast.435.clone.1, f32[3,4096,48,443]{3,2,1,0} %gather.16.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.53.clone.1 = f32[3,4096,48,443]{3,2,1,0} add(f32[3,4096,48,443]{3,2,1,0} %multiply.52.clone.1, f32[3,4096,48,443]{3,2,1,0} %multiply.51.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.84.clone.1 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.89.clone.1, f32[3,4096,443]{2,1,0} %subtract.86.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.387.clone.1 = f32[3,4096,48,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.84.clone.1), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.45.clone.1 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %add.146.clone.1 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.72.clone.1, s32[3,4096,443]{2,1,0} %broadcast.274), metadata={op_name="jit(interpolate_backward)/jit(main)/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.51.clone.1 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.279, s32[3,4096,443]{2,1,0} %add.146.clone.1, s32[3,4096,443]{2,1,0} %broadcast.278), metadata={op_name="jit(interpolate_backward)/jit(main)/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.69.clone.1 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.51.clone.1, s32[3,4096,443]{2,1,0} %broadcast.279), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.145.clone.1 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.51.clone.1, s32[3,4096,443]{2,1,0} %broadcast.277), metadata={op_name="jit(interpolate_backward)/jit(main)/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.69.clone.1 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.69.clone.1, s32[3,4096,443]{2,1,0} %add.145.clone.1, s32[3,4096,443]{2,1,0} %clamp.51.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.87.clone.1 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.69.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.47.clone.1 = s32[3,4096,443,3]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.45.clone.1, s32[3,4096,443,1]{2,1,0,3} %bitcast.87.clone.1, s32[3,4096,443,1]{2,1,0,3} %bitcast.102.clone.1), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.15.clone.1 = f32[3,4096,48,443]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.224, s32[3,4096,443,3]{2,1,0,3} %concatenate.47.clone.1), 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(interpolate_backward)/jit(main)/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.50.clone.1 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %broadcast.387.clone.1, f32[3,4096,48,443]{3,2,1,0} %gather.15.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.52.clone.1 = f32[3,4096,48,443]{3,2,1,0} add(f32[3,4096,48,443]{3,2,1,0} %add.53.clone.1, f32[3,4096,48,443]{3,2,1,0} %multiply.50.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.80.clone.1 = f32[3,4096,443]{2,1,0} multiply(f32[3,4096,443]{2,1,0} %subtract.89.clone.1, f32[3,4096,443]{2,1,0} %subtract.87.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.304.clone.1 = f32[3,4096,48,443]{3,2,1,0} broadcast(f32[3,4096,443]{2,1,0} %multiply.80.clone.1), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.41.clone.1 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.75.clone.1 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.69.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.43.clone.1 = s32[3,4096,443,3]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.41.clone.1, s32[3,4096,443,1]{3,2,1,0} %bitcast.75.clone.1, s32[3,4096,443,1]{3,2,1,0} %bitcast.94.clone.1), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.14.clone.1 = f32[3,4096,48,443]{3,2,1,0} gather(f32[3,48,128,128]{3,2,1,0} %param_2.224, s32[3,4096,443,3]{3,2,1,0} %concatenate.43.clone.1), 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(interpolate_backward)/jit(main)/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.49.clone.1 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %broadcast.304.clone.1, f32[3,4096,48,443]{3,2,1,0} %gather.14.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.51.clone.1 = f32[3,4096,48,443]{3,2,1,0} add(f32[3,4096,48,443]{3,2,1,0} %add.52.clone.1, f32[3,4096,48,443]{3,2,1,0} %multiply.49.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.48 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %copy.14, f32[3,4096,48,443]{3,2,1,0} %add.51.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
  %exponential.2 = f32[3,4096,48,443]{3,2,1,0} exponential(f32[3,4096,48,443]{3,2,1,0} %multiply.48), metadata={op_name="jit(interpolate_backward)/jit(main)/exp" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=47}
  ROOT %tuple.6 = (f32[3,4096,48,443]{3,2,1,0}, f32[3,4096,48,443]{3,2,1,0}) tuple(f32[3,4096,48,443]{3,2,1,0} %exponential.2, f32[3,4096,48,443]{3,2,1,0} %add.51.clone.1)
}

%fused_computation.16 (param_0.107: f32[3,4096,443], param_1.132: s32[3,1], param_2.119: s32[3,1]) -> (f32[3,4096,443], f32[3,4096,443], f32[3,4096,443]) {
  %param_0.107 = f32[3,4096,443]{2,1,0} parameter(0)
  %bitcast.51 = f32[1,3,4096,443]{3,2,0,1} bitcast(f32[3,4096,443]{2,1,0} %param_0.107), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 443) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %param_2.119 = s32[3,1]{1,0} parameter(2)
  %gather.33 = f32[3,4096,443]{2,1,0} gather(f32[3,4096,443]{2,1,0} %param_0.107, s32[3,1]{1,0} %param_2.119), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,4096,443}, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443) 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.50 = f32[1,3,4096,443]{3,2,0,1} bitcast(f32[3,4096,443]{2,1,0} %gather.33), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 443) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %param_1.132 = s32[3,1]{1,0} parameter(1)
  %gather.32 = f32[3,4096,443]{2,1,0} gather(f32[3,4096,443]{2,1,0} %param_0.107, s32[3,1]{1,0} %param_1.132), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,4096,443}, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443) 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.49 = f32[1,3,4096,443]{3,2,0,1} bitcast(f32[3,4096,443]{2,1,0} %gather.32), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/broadcast_in_dim[shape=(1, 3, 4096, 443) broadcast_dimensions=(1, 2, 3)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %concatenate.25 = f32[3,3,4096,443]{3,2,0,1} concatenate(f32[1,3,4096,443]{3,2,0,1} %bitcast.51, f32[1,3,4096,443]{3,2,0,1} %bitcast.50, f32[1,3,4096,443]{3,2,0,1} %bitcast.49), dimensions={0}, metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/concatenate[dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=60}
  %constant_101 = f32[] constant(1)
  %broadcast.222 = f32[3,3,4096,443]{3,2,0,1} broadcast(f32[] %constant_101), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=159}
  %add.72 = f32[3,3,4096,443]{3,2,0,1} add(f32[3,3,4096,443]{3,2,0,1} %concatenate.25, f32[3,3,4096,443]{3,2,0,1} %broadcast.222), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/add" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=159}
  %constant_100 = f32[] constant(63.5)
  %broadcast.220 = f32[3,3,4096,443]{3,2,0,1} broadcast(f32[] %constant_100), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=162}
  %multiply.68 = f32[3,3,4096,443]{3,2,0,1} multiply(f32[3,3,4096,443]{3,2,0,1} %add.72, f32[3,3,4096,443]{3,2,0,1} %broadcast.220), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(jvp(interpolate))/jit(jit_jvp(interpolate))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=162}
  %transpose.13 = f32[3,4096,3,443]{3,1,0,2} transpose(f32[3,3,4096,443]{3,2,0,1} %multiply.68), dimensions={0,2,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.9 = f32[3,4096,1,443]{3,1,0,2} slice(f32[3,4096,3,443]{3,1,0,2} %transpose.13), slice={[0:3], [0:4096], [1:2], [0:443]}, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.27 = f32[3,4096,443]{2,1,0} bitcast(f32[3,4096,1,443]{3,1,0,2} %slice.9), metadata={op_name="jit(interpolate_backward)/jit(main)/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.10.clone.1 = f32[3,4096,1,443]{3,1,0,2} slice(f32[3,4096,3,443]{3,1,0,2} %transpose.13), slice={[0:3], [0:4096], [0:1], [0:443]}, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443) 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.30.clone.1 = f32[3,4096,443]{2,1,0} bitcast(f32[3,4096,1,443]{3,1,0,2} %slice.10.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.8.clone.1 = f32[3,4096,1,443]{3,1,0,2} slice(f32[3,4096,3,443]{3,1,0,2} %transpose.13), slice={[0:3], [0:4096], [2:3], [0:443]}, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443) strides=(1, 1, 1, 1)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %bitcast.26.clone.1 = f32[3,4096,443]{2,1,0} bitcast(f32[3,4096,1,443]{3,1,0,2} %slice.8.clone.1), metadata={op_name="jit(interpolate_backward)/jit(main)/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.8 = (f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}) tuple(f32[3,4096,443]{2,1,0} %bitcast.27, f32[3,4096,443]{2,1,0} %bitcast.30.clone.1, f32[3,4096,443]{2,1,0} %bitcast.26.clone.1)
}

%fused_computation.23 (param_0.158: f32[3,48,128], param_1.217: f32[3,4096,443], param_2.218: f32[3,4096,48,443], param_3.184: f32[]) -> f32[3,4096,48,443] {
  %constant_179 = f32[] constant(1)
  %broadcast.302 = f32[3,4096,443]{2,1,0} broadcast(f32[] %constant_179), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.217 = f32[3,4096,443]{2,1,0} parameter(1)
  %floor.51 = f32[3,4096,443]{2,1,0} floor(f32[3,4096,443]{2,1,0} %param_1.217), metadata={op_name="jit(interpolate_backward)/jit(main)/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.49 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %param_1.217, f32[3,4096,443]{2,1,0} %floor.51), metadata={op_name="jit(interpolate_backward)/jit(main)/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.48 = f32[3,4096,443]{2,1,0} subtract(f32[3,4096,443]{2,1,0} %broadcast.302, f32[3,4096,443]{2,1,0} %subtract.49), metadata={op_name="jit(interpolate_backward)/jit(main)/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.301 = f32[3,4096,48,443]{2,1,0,3} broadcast(f32[3,4096,443]{2,1,0} %subtract.48), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.158 = f32[3,48,128]{2,1,0} parameter(0)
  %iota.39 = s32[3,4096,443,1]{3,2,1,0} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_175 = s32[] constant(0)
  %broadcast.300 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_175), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.27 = s32[3,4096,443]{2,1,0} convert(f32[3,4096,443]{2,1,0} %floor.51), metadata={op_name="jit(interpolate_backward)/jit(main)/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_173 = s32[] constant(127)
  %broadcast.299 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_173), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.300, s32[3,4096,443]{2,1,0} %convert.27, s32[3,4096,443]{2,1,0} %broadcast.299), metadata={op_name="jit(interpolate_backward)/jit(main)/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.43 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.27, s32[3,4096,443]{2,1,0} %broadcast.300), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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_172 = s32[] constant(128)
  %broadcast.298 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_172), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.100 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.27, s32[3,4096,443]{2,1,0} %broadcast.298), metadata={op_name="jit(interpolate_backward)/jit(main)/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.43 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.43, s32[3,4096,443]{2,1,0} %add.100, s32[3,4096,443]{2,1,0} %clamp.27), metadata={op_name="jit(interpolate_backward)/jit(main)/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.67 = s32[3,4096,443,1]{3,2,1,0} bitcast(s32[3,4096,443]{2,1,0} %select.43), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.41 = s32[3,4096,443,2]{3,2,1,0} concatenate(s32[3,4096,443,1]{3,2,1,0} %iota.39, s32[3,4096,443,1]{3,2,1,0} %bitcast.67), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.41 = f32[3,4096,48,443]{2,1,0,3} gather(f32[3,48,128]{2,1,0} %param_0.158, s32[3,4096,443,2]{3,2,1,0} %concatenate.41), 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(interpolate_backward)/jit(main)/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.78 = f32[3,4096,48,443]{2,1,0,3} multiply(f32[3,4096,48,443]{2,1,0,3} %broadcast.301, f32[3,4096,48,443]{2,1,0,3} %gather.41), metadata={op_name="jit(interpolate_backward)/jit(main)/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.297 = f32[3,4096,48,443]{2,1,0,3} broadcast(f32[3,4096,443]{2,1,0} %subtract.49), dimensions={0,1,3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %iota.38 = s32[3,4096,443,1]{2,1,0,3} iota(), iota_dimension=0, metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) dimension=0]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %constant_174 = s32[] constant(1)
  %broadcast.295 = s32[3,4096,443]{2,1,0} broadcast(s32[] %constant_174), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.99 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %convert.27, s32[3,4096,443]{2,1,0} %broadcast.295), metadata={op_name="jit(interpolate_backward)/jit(main)/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.26 = s32[3,4096,443]{2,1,0} clamp(s32[3,4096,443]{2,1,0} %broadcast.300, s32[3,4096,443]{2,1,0} %add.99, s32[3,4096,443]{2,1,0} %broadcast.299), metadata={op_name="jit(interpolate_backward)/jit(main)/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.42 = pred[3,4096,443]{2,1,0} compare(s32[3,4096,443]{2,1,0} %clamp.26, s32[3,4096,443]{2,1,0} %broadcast.300), direction=LT, metadata={op_name="jit(interpolate_backward)/jit(main)/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.98 = s32[3,4096,443]{2,1,0} add(s32[3,4096,443]{2,1,0} %clamp.26, s32[3,4096,443]{2,1,0} %broadcast.298), metadata={op_name="jit(interpolate_backward)/jit(main)/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.42 = s32[3,4096,443]{2,1,0} select(pred[3,4096,443]{2,1,0} %compare.42, s32[3,4096,443]{2,1,0} %add.98, s32[3,4096,443]{2,1,0} %clamp.26), metadata={op_name="jit(interpolate_backward)/jit(main)/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.66 = s32[3,4096,443,1]{2,1,0,3} bitcast(s32[3,4096,443]{2,1,0} %select.42), metadata={op_name="jit(interpolate_backward)/jit(main)/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, 443, 1) broadcast_dimensions=(0, 1, 2)]" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=222}
  %concatenate.40 = s32[3,4096,443,2]{2,1,0,3} concatenate(s32[3,4096,443,1]{2,1,0,3} %iota.38, s32[3,4096,443,1]{2,1,0,3} %bitcast.66), dimensions={3}, metadata={op_name="jit(interpolate_backward)/jit(main)/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.40 = f32[3,4096,48,443]{2,1,0,3} gather(f32[3,48,128]{2,1,0} %param_0.158, s32[3,4096,443,2]{2,1,0,3} %concatenate.40), 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(interpolate_backward)/jit(main)/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.77 = f32[3,4096,48,443]{2,1,0,3} multiply(f32[3,4096,48,443]{2,1,0,3} %broadcast.297, f32[3,4096,48,443]{2,1,0,3} %gather.40), metadata={op_name="jit(interpolate_backward)/jit(main)/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.97 = f32[3,4096,48,443]{2,1,0,3} add(f32[3,4096,48,443]{2,1,0,3} %multiply.78, f32[3,4096,48,443]{2,1,0,3} %multiply.77), metadata={op_name="jit(interpolate_backward)/jit(main)/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.16 = f32[3,4096,48,443]{3,2,1,0} copy(f32[3,4096,48,443]{2,1,0,3} %add.97), metadata={op_name="jit(interpolate_backward)/jit(main)/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_3.184 = f32[] parameter(3)
  %divide.3 = f32[] divide(f32[] %constant_179, f32[] %param_3.184), metadata={op_name="jit(interpolate_backward)/jit(main)/div" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=45}
  %broadcast.514 = f32[3,4096,48,443]{3,2,1,0} broadcast(f32[] %divide.3), dimensions={}, metadata={op_name="jit(interpolate_backward)/jit(main)/broadcast_in_dim[shape=(144, 4096, 443) broadcast_dimensions=()]" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=46}
  %param_2.218 = f32[3,4096,48,443]{3,2,1,0} parameter(2)
  %multiply.96 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %broadcast.514, f32[3,4096,48,443]{3,2,1,0} %param_2.218), metadata={op_name="jit(interpolate_backward)/jit(main)/mul" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=47}
  ROOT %multiply.74 = f32[3,4096,48,443]{3,2,1,0} multiply(f32[3,4096,48,443]{3,2,1,0} %copy.16, f32[3,4096,48,443]{3,2,1,0} %multiply.96), metadata={op_name="jit(interpolate_backward)/jit(main)/jit(transpose(jvp(interpolate)))/jit(jit_transpose(jvp(interpolate)))/mul" source_file="/home/brent/tensorf-jax/tensorf/tensor_vm.py" source_line=176}
}

ENTRY %main.399 (Arg_0.1: f32[3,48,128], Arg_1.2: f32[3,48,128,128], Arg_2.3: f32[3,4096,443]) -> (f32[3,48,128], f32[3,48,128,128]) {
  %Arg_0.1 = f32[3,48,128]{2,1,0} parameter(0)
  %Arg_2.3 = f32[3,4096,443]{2,1,0} parameter(2)
  %constant_51 = s32[3,1]{1,0} constant({ {1}, {2}, {0} }), metadata={op_name="jit(interpolate_backward)/jit(main)/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_47 = s32[3,1]{1,0} constant({ {2}, {0}, {1} }), metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %fusion.16 = (f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}) fusion(f32[3,4096,443]{2,1,0} %Arg_2.3, s32[3,1]{1,0} %constant_51, s32[3,1]{1,0} %constant_47), kind=kLoop, calls=%fused_computation.16, metadata={op_name="jit(interpolate_backward)/jit(main)/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.17 = f32[3,4096,443]{2,1,0} get-tuple-element((f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}) %fusion.16), index=1
  %Arg_1.2 = f32[3,48,128,128]{3,2,1,0} parameter(1)
  %get-tuple-element.18 = f32[3,4096,443]{2,1,0} get-tuple-element((f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}) %fusion.16), index=2
  %get-tuple-element.16 = f32[3,4096,443]{2,1,0} get-tuple-element((f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}, f32[3,4096,443]{2,1,0}) %fusion.16), index=0
  %fusion.1 = (f32[3,4096,48,443]{3,2,1,0}, f32[3,4096,48,443]{3,2,1,0}) fusion(f32[3,48,128]{2,1,0} %Arg_0.1, f32[3,4096,443]{2,1,0} %get-tuple-element.17, f32[3,48,128,128]{3,2,1,0} %Arg_1.2, f32[3,4096,443]{2,1,0} %get-tuple-element.18, f32[3,4096,443]{2,1,0} %get-tuple-element.16), kind=kLoop, calls=%fused_computation.1, metadata={op_name="jit(interpolate_backward)/jit(main)/exp" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=47}
  %get-tuple-element.15 = f32[3,4096,48,443]{3,2,1,0} get-tuple-element((f32[3,4096,48,443]{3,2,1,0}, f32[3,4096,48,443]{3,2,1,0}) %fusion.1), index=1
  %get-tuple-element.14 = f32[3,4096,48,443]{3,2,1,0} get-tuple-element((f32[3,4096,48,443]{3,2,1,0}, f32[3,4096,48,443]{3,2,1,0}) %fusion.1), index=0
  %input_fusion_reduce = f32[16165]{0} fusion(f32[3,4096,48,443]{3,2,1,0} %get-tuple-element.14), kind=kInput, calls=%input_fused_computation_reduce
  %constant_5 = f32[] constant(0)
  %reduce.290 = f32[] reduce(f32[16165]{0} %input_fusion_reduce, f32[] %constant_5), dimensions={0}, to_apply=%region_0.286, metadata={op_name="jit(interpolate_backward)/jit(main)/reduce_sum[axes=(0, 1, 2)]" source_file="time_vmap_and_generate_hlo_just_interpolate.py" source_line=46}
  %input_fusion_scatter.5 = f32[3,48,128]{2,1,0} fusion(f32[3,4096,48,443]{3,2,1,0} %get-tuple-element.15, f32[3,4096,443]{2,1,0} %get-tuple-element.17, f32[3,4096,48,443]{3,2,1,0} %get-tuple-element.14, f32[] %reduce.290), kind=kInput, calls=%input_fused_computation_scatter.5, metadata={op_name="jit(interpolate_backward)/jit(main)/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,48,128]{2,1,0} fusion(f32[3,4096,48,443]{3,2,1,0} %get-tuple-element.15, f32[3,4096,443]{2,1,0} %get-tuple-element.17, f32[3,4096,48,443]{3,2,1,0} %get-tuple-element.14, f32[] %reduce.290), kind=kInput, calls=%input_fused_computation_scatter.4, metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %add.41 = f32[3,48,128]{2,1,0} add(f32[3,48,128]{2,1,0} %input_fusion_scatter.5, f32[3,48,128]{2,1,0} %input_fusion_scatter.4), metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %fusion.23 = f32[3,4096,48,443]{3,2,1,0} fusion(f32[3,48,128]{2,1,0} %Arg_0.1, f32[3,4096,443]{2,1,0} %get-tuple-element.17, f32[3,4096,48,443]{3,2,1,0} %get-tuple-element.14, f32[] %reduce.290), kind=kLoop, calls=%fused_computation.23, metadata={op_name="jit(interpolate_backward)/jit(main)/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,48,128,128]{3,2,1,0} fusion(f32[3,4096,48,443]{3,2,1,0} %fusion.23, f32[3,4096,443]{2,1,0} %get-tuple-element.18, f32[3,4096,443]{2,1,0} %get-tuple-element.16), kind=kInput, calls=%input_fused_computation_scatter.3, metadata={op_name="jit(interpolate_backward)/jit(main)/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,48,128,128]{3,2,1,0} fusion(f32[3,4096,48,443]{3,2,1,0} %fusion.23, f32[3,4096,443]{2,1,0} %get-tuple-element.18, f32[3,4096,443]{2,1,0} %get-tuple-element.16), kind=kInput, calls=%input_fused_computation_scatter.2, metadata={op_name="jit(interpolate_backward)/jit(main)/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,48,128,128]{3,2,1,0} fusion(f32[3,4096,48,443]{3,2,1,0} %fusion.23, f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.3, f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.2, f32[3,4096,443]{2,1,0} %get-tuple-element.18, f32[3,4096,443]{2,1,0} %get-tuple-element.16), kind=kInput, calls=%input_fused_computation_scatter.1, metadata={op_name="jit(interpolate_backward)/jit(main)/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,48,128,128]{3,2,1,0} fusion(f32[3,4096,48,443]{3,2,1,0} %fusion.23, f32[3,4096,443]{2,1,0} %get-tuple-element.18, f32[3,4096,443]{2,1,0} %get-tuple-element.16), kind=kInput, calls=%input_fused_computation_scatter, metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  %add.45 = f32[3,48,128,128]{3,2,1,0} add(f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter.1, f32[3,48,128,128]{3,2,1,0} %input_fusion_scatter), metadata={op_name="jit(interpolate_backward)/jit(main)/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}
  ROOT %tuple.398 = (f32[3,48,128]{2,1,0}, f32[3,48,128,128]{3,2,1,0}) tuple(f32[3,48,128]{2,1,0} %add.41, f32[3,48,128,128]{3,2,1,0} %add.45)
}