Fast decomposable/factorized convolutions

[helange23]

I need to convolve 3D data with a decomposable kernel K of shape (35,35,6,6,6) with 35 being the channel dimension.
I can decompose this kernel into 3 kernels Ki of shapes (35,35,6,1,1), (35,35,1,6,1) and (35,35,1,1,6) such that K1*K2*K3 = K with broadcasting.

In principle, this should lead to a 12x reduction in FLOPs [(6*6*6)/(6*3)] but it actually leads to a 3x slow down. Below is an example of what I am doing.

I would really like to not pay the price for a full convolution when I am able to decompose it. Has anyone gotten something like this to work?

Also, wrong_fac_conv and fac_conv technically require the amount of FLOPs. One is however 8x faster than the other.
Would writing my own custom_call be the only option to get the FLOPs out of JAX?

import jax.numpy as jnp
import numpy as np
import jax

kx = jnp.array(np.random.uniform(-1,1,(35,35,6,1,1)))
ky = jnp.array(np.random.uniform(-1,1,(35,35,1,6,1)))
kz = jnp.array(np.random.uniform(-1,1,(35,35,1,1,6)))

kernel = kx*ky*kz

M = jnp.array(np.random.uniform(-1,1,(1,128,128,128,35)))
M2 = M.transpose((4,0,1,2,3))[:,0]

dn = jax.lax.conv_dimension_numbers(M.shape, kernel.shape,
                                    ('NHWDC', 'OIHWD', 'NHWDC'))

@jax.jit
def reg_conv(M, kernel):
    return jax.lax.conv_general_dilated(M,    # lhs = image tensor
                                        kernel,  # rhs = conv kernel tensor
                                        (1,1,1), # window strides
                                        "SAME",  # padding mode
                                        (1,1,1), # lhs/image dilation
                                        (1,1,1), # rhs/kernel dilation
                                        dn)     # dimension_numbers

def factored_dim(kernel_x, kernel_y, kernel_z, M):
    carry = (jnp.zeros((1,128,128,128,1)))
    def scan_fun(carry, x):
        M, kx, ky, kz = x
        L_ = carry
        L = reg_conv(M[None,...,None], kz[None,None])
        L = reg_conv(L, ky[None,None])
        L = reg_conv(L, kx[None,None])
        L_ += L
        return L_, None
    xs = (M, kernel_x, kernel_y, kernel_z)
    out, _ = jax.lax.scan(scan_fun, carry, xs)
    return out[...,0]

@jax.jit
def fac_conv(M, kernel_x, kernel_y, kernel_z):
    f = lambda x,y,z: factored_dim(x,y,z, M)
    return jax.vmap(f, out_axes=-1)(kernel_x, kernel_y, kernel_z)


@jax.jit
def wrong_fac_conv(M, kx, ky, kz):
    L = reg_conv(M, kz)
    L = reg_conv(L, ky)
    L = reg_conv(L, kx)
    return L

o = reg_conv(M, kernel)
o2 = fac_conv(M2, kx, ky, kz)
o3 = wrong_fac_conv(M, kx, ky, kz)

print('err', jnp.mean(jnp.abs(o-o2)))
print('err wrong', jnp.mean(jnp.abs(o-o3)))

>err 2.0008392e-06
>err wrong 267.67865

%timeit reg_conv(M, kernel).block_until_ready()
144 ms ± 346 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit fac_conv(M2, kx, ky, kz).block_until_ready()
393 ms ± 718 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit wrong_fac_conv(M, kx, ky, kz).block_until_ready()
48.6 ms ± 349 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

[hawkinsp]

One key piece of information that would help: what hardware are you using and what software versions are you using? It's hard to say what's happening here without running it and looking at a profile.

A tensorboard profile would be interesting if you wanted to try grabbing one: https://jax.readthedocs.io/en/latest/profiling.html

[helange23]

I am using jax version 0.2.7, jaxlib 0.1.57 and CUDA 11 on a Gefore RTX 2080Ti.
Does that mean that you cannot reproduce this behavior? Does fac_conv run comparably as fast as wrong_fac_conv on your setup?

For some reason, I cannot get Tensorboard profiling to work but attached is the hopefully relevant part from the XLA profile:

2021-06-25 16:08:31.360352: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] ********** categories table for microseconds above estimated optimum **********
2021-06-25 16:08:31.360364: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.360375: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]  295 (56.44% Σ56.44%)   pad (2 ops)
2021-06-25 16:08:31.360387: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 28.40% %pad.3 = f32[1,128,133,128,35]{3,2,1,4,0} pad(f32[1,128,128,128,35]{3,2,1,4,0}, f32[]), padding=0_0x0_0x2_3x0_0x0_0
2021-06-25 16:08:31.360399: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 28.04% %pad.4 = f32[1,133,128,128,35]{3,2,1,4,0} pad(f32[1,128,128,128,35]{3,2,1,4,0}, f32[]), padding=0_0x2_3x0_0x0_0x0_0
2021-06-25 16:08:31.360411: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]  219 (41.99% Σ98.43%)   loop fusion (5 ops)
2021-06-25 16:08:31.360431: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 37.56% %fusion = f32[35,1,128,128,128,1]{4,3,2,0,5,1} fusion(f32[35,1,128,128,128,1]{4,3,2,0,5,1}, f32[1,128,128,128,35]{3,2,1,4,0}), kind=kLoop, calls=%fused_computation, metadata={op_type="add" op_name="jit(fac_conv)/scan/while/body/add" source_file="decomposable_conv.py" source_line=48}
2021-06-25 16:08:31.360444: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  1.73% %fusion.4 = f32[1,128,128,133,1]{3,2,1,4,0} fusion(f32[35,128,128,128]{3,2,1,0}, s32[]), kind=kLoop, calls=%fused_computation.4
2021-06-25 16:08:31.360456: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  1.00% %fusion.1 = f32[35,1,6,1,1]{4,3,2,1,0} fusion(f32[35,35,6,1,1]{2,1,4,3,0}, s32[]), kind=kLoop, calls=%fused_computation.1, metadata={op_type="reshape" op_name="jit(fac_conv)/scan/while/body/jit(vmap(reg_conv))/reshape[ dimensions=None\n                                                           new_sizes=(35, 1, 6, 1, 1) ]" source_file="decomposable_conv.py" source_line=38}
2021-06-25 16:08:31.360468: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.85% %fusion.3 = f32[35,1,1,1,6]{4,3,2,1,0} fusion(f32[35,35,1,1,6]{4,1,3,2,0}, s32[]), kind=kLoop, calls=%fused_computation.3, metadata={op_type="reshape" op_name="jit(fac_conv)/scan/while/body/jit(vmap(reg_conv))/reshape[ dimensions=None\n                                                           new_sizes=(35, 1, 1, 1, 6) ]" source_file="decomposable_conv.py" source_line=38}
2021-06-25 16:08:31.360480: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.84% %fusion.2 = f32[35,1,1,6,1]{4,3,2,1,0} fusion(f32[35,35,1,6,1]{3,1,4,2,0}, s32[]), kind=kLoop, calls=%fused_computation.2, metadata={op_type="reshape" op_name="jit(fac_conv)/scan/while/body/jit(vmap(reg_conv))/reshape[ dimensions=None\n                                                           new_sizes=(35, 1, 1, 6, 1) ]" source_file="decomposable_conv.py" source_line=38}
2021-06-25 16:08:31.360492: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]    5 ( 1.05% Σ99.47%)   non-fusion elementwise (1 ops)
2021-06-25 16:08:31.360504: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  1.05% %add.175 = s32[] add(s32[], s32[]), control-predecessors={%fusion.1, %fusion.4, %fusion.3, %fusion.2}, metadata={op_type="add" op_name="jit(fac_conv)/scan/while/body/add" source_file="decomposable_conv.py" source_line=68}
2021-06-25 16:08:31.360516: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]    3 ( 0.53% Σ100.00%)   tuple (1 ops)
2021-06-25 16:08:31.360528: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.53% %tuple.13 = (f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, s32[], f32[35,1,128,128,128,1]{4,3,2,0,5,1}) tuple(f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, ...(+2))
2021-06-25 16:08:31.360540: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.360551: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.360562: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] ********** microseconds report **********
2021-06-25 16:08:31.360574: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 11,423 microseconds in total.
2021-06-25 16:08:31.360586: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 79 microseconds ( 0.69%) not accounted for by the data.
2021-06-25 16:08:31.360598: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 12 ops.
2021-06-25 16:08:31.360614: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.360626: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] ********** categories table for microseconds **********
2021-06-25 16:08:31.360637: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.360649: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]   7,420 (64.96% Σ65.65%)   custom-call (3 ops)
2021-06-25 16:08:31.360663: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 28.72% %custom-call.8 = (f32[1,128,128,128,35]{3,2,1,4,0}, u8[0]{0}) custom-call(f32[1,133,128,128,35]{3,2,1,4,0}, f32[35,1,6,1,1]{4,3,2,1,0}), window={size=6x1x1}, dim_labels=b012f_oi012->b012f, feature_group_count=35, custom_call_target="__cudnn$convForward", metadata={op_type="conv_general_dilated" op_name="jit(fac_conv)/scan/while/body/jit(vmap(reg_conv))/conv_general_dilated[ batch_group_count=1\n                                                                        dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 4, 1, 2, 3), rhs_spec=(0, 1, 2, 3, 4), out_spec=(0, 4, 1, 2, 3))\n                                                                        feature_group_count=35\n                                                                        lhs_dilation=(1, 1, 1)\n                                                                        lhs_shape=(1, 128, 128, 128, 35)\n                                                                        padding=((2, 3), (0, 0), (0, 0))\n                                                                        precision=None\n                                                                        rhs_dilation=(1, 1, 1)\n                                                                        rhs_shape=(35, 1, 6, 1, 1)\n                                                                        window_strides=(1, 1, 1) ]" source_file="decomposable_conv.py" source_line=38}, backend_config="{\"algorithm\":\"1\",\"tensor_ops_enabled\":true,\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"
2021-06-25 16:08:31.360678: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 26.98% %custom-call.7 = (f32[1,128,128,128,35]{3,2,1,4,0}, u8[0]{0}) custom-call(f32[1,128,133,128,35]{3,2,1,4,0}, f32[35,1,1,6,1]{4,3,2,1,0}), window={size=1x6x1}, dim_labels=b012f_oi012->b012f, feature_group_count=35, custom_call_target="__cudnn$convForward", metadata={op_type="conv_general_dilated" op_name="jit(fac_conv)/scan/while/body/jit(vmap(reg_conv))/conv_general_dilated[ batch_group_count=1\n                                                                        dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 4, 1, 2, 3), rhs_spec=(0, 1, 2, 3, 4), out_spec=(0, 4, 1, 2, 3))\n                                                                        feature_group_count=35\n                                                                        lhs_dilation=(1, 1, 1)\n                                                                        lhs_shape=(1, 128, 128, 128, 35)\n                                                                        padding=((0, 0), (2, 3), (0, 0))\n                                                                        precision=None\n                                                                        rhs_dilation=(1, 1, 1)\n                                                                        rhs_shape=(35, 1, 1, 6, 1)\n                                                                        window_strides=(1, 1, 1) ]" source_file="decomposable_conv.py" source_line=38}, backend_config="{\"algorithm\":\"1\",\"tensor_ops_enabled\":false,\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"
2021-06-25 16:08:31.360692: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  9.26% %custom-call.6 = (f32[1,128,128,128,35]{3,2,1,4,0}, u8[0]{0}) custom-call(f32[1,128,128,133,1]{3,2,1,4,0}, f32[35,1,1,1,6]{4,3,2,1,0}), window={size=1x1x6}, dim_labels=b012f_oi012->b012f, custom_call_target="__cudnn$convForward", metadata={op_type="conv_general_dilated" op_name="jit(fac_conv)/scan/while/body/jit(vmap(reg_conv))/conv_general_dilated[ batch_group_count=1\n                                                                        dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 4, 1, 2, 3), rhs_spec=(0, 1, 2, 3, 4), out_spec=(0, 4, 1, 2, 3))\n                                                                        feature_group_count=1\n                                                                        lhs_dilation=(1, 1, 1)\n                                                                        lhs_shape=(1, 128, 128, 128, 1)\n                                                                        padding=((0, 0), (0, 0), (2, 3))\n                                                                        precision=None\n                                                                        rhs_dilation=(1, 1, 1)\n                                                                        rhs_shape=(35, 1, 1, 1, 6)\n                                                                        window_strides=(1, 1, 1) ]" source_file="decomposable_conv.py" source_line=38}, backend_config="{\"algorithm\":\"1\",\"tensor_ops_enabled\":false,\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"
2021-06-25 16:08:31.360709: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]   2,239 (19.60% Σ85.25%)   pad (2 ops)
2021-06-25 16:08:31.360721: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  9.81% %pad.3 = f32[1,128,133,128,35]{3,2,1,4,0} pad(f32[1,128,128,128,35]{3,2,1,4,0}, f32[]), padding=0_0x0_0x2_3x0_0x0_0
2021-06-25 16:08:31.360733: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  9.79% %pad.4 = f32[1,133,128,128,35]{3,2,1,4,0} pad(f32[1,128,128,128,35]{3,2,1,4,0}, f32[]), padding=0_0x2_3x0_0x0_0x0_0
2021-06-25 16:08:31.360744: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]   1,677 (14.68% Σ99.93%)   loop fusion (5 ops)
2021-06-25 16:08:31.360756: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 14.24% %fusion = f32[35,1,128,128,128,1]{4,3,2,0,5,1} fusion(f32[35,1,128,128,128,1]{4,3,2,0,5,1}, f32[1,128,128,128,35]{3,2,1,4,0}), kind=kLoop, calls=%fused_computation, metadata={op_type="add" op_name="jit(fac_conv)/scan/while/body/add" source_file="decomposable_conv.py" source_line=48}
2021-06-25 16:08:31.360768: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.32% %fusion.4 = f32[1,128,128,133,1]{3,2,1,4,0} fusion(f32[35,128,128,128]{3,2,1,0}, s32[]), kind=kLoop, calls=%fused_computation.4
2021-06-25 16:08:31.360780: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.05% %fusion.1 = f32[35,1,6,1,1]{4,3,2,1,0} fusion(f32[35,35,6,1,1]{2,1,4,3,0}, s32[]), kind=kLoop, calls=%fused_computation.1, metadata={op_type="reshape" op_name="jit(fac_conv)/scan/while/body/jit(vmap(reg_conv))/reshape[ dimensions=None\n                                                           new_sizes=(35, 1, 6, 1, 1) ]" source_file="decomposable_conv.py" source_line=38}
2021-06-25 16:08:31.360792: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.04% %fusion.3 = f32[35,1,1,1,6]{4,3,2,1,0} fusion(f32[35,35,1,1,6]{4,1,3,2,0}, s32[]), kind=kLoop, calls=%fused_computation.3, metadata={op_type="reshape" op_name="jit(fac_conv)/scan/while/body/jit(vmap(reg_conv))/reshape[ dimensions=None\n                                                           new_sizes=(35, 1, 1, 1, 6) ]" source_file="decomposable_conv.py" source_line=38}
2021-06-25 16:08:31.360808: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.04% %fusion.2 = f32[35,1,1,6,1]{4,3,2,1,0} fusion(f32[35,35,1,6,1]{3,1,4,2,0}, s32[]), kind=kLoop, calls=%fused_computation.2, metadata={op_type="reshape" op_name="jit(fac_conv)/scan/while/body/jit(vmap(reg_conv))/reshape[ dimensions=None\n                                                           new_sizes=(35, 1, 1, 6, 1) ]" source_file="decomposable_conv.py" source_line=38}
2021-06-25 16:08:31.360820: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]       5 ( 0.05% Σ99.98%)   non-fusion elementwise (1 ops)
2021-06-25 16:08:31.360832: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.05% %add.175 = s32[] add(s32[], s32[]), control-predecessors={%fusion.1, %fusion.4, %fusion.3, %fusion.2}, metadata={op_type="add" op_name="jit(fac_conv)/scan/while/body/add" source_file="decomposable_conv.py" source_line=68}
2021-06-25 16:08:31.360844: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]       3 ( 0.02% Σ100.00%)   tuple (1 ops)
2021-06-25 16:08:31.360855: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.02% %tuple.13 = (f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, s32[], f32[35,1,128,128,128,1]{4,3,2,0,5,1}) tuple(f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, ...(+2))
2021-06-25 16:08:31.360867: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.360878: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.360890: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] ********** MiB read+written report **********
2021-06-25 16:08:31.360901: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 1,998 MiB read+written in total.
2021-06-25 16:08:31.360933: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 0 MiB read+written ( 0.00%) not accounted for by the data.
2021-06-25 16:08:31.360944: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 24 ops.
2021-06-25 16:08:31.360956: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.360968: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] ********** categories table for MiB read+written **********
2021-06-25 16:08:31.360979: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.361010: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]  1,142 (57.15% Σ57.15%)   pad (2 ops)
2021-06-25 16:08:31.361071: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 28.57% %pad.3 = f32[1,128,133,128,35]{3,2,1,4,0} pad(f32[1,128,128,128,35]{3,2,1,4,0}, f32[]), padding=0_0x0_0x2_3x0_0x0_0
2021-06-25 16:08:31.361086: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 28.57% %pad.4 = f32[1,133,128,128,35]{3,2,1,4,0} pad(f32[1,128,128,128,35]{3,2,1,4,0}, f32[]), padding=0_0x2_3x0_0x0_0x0_0
2021-06-25 16:08:31.361099: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]    856 (42.85% Σ100.00%)   loop fusion (5 ops)
2021-06-25 16:08:31.361111: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 42.04% %fusion = f32[35,1,128,128,128,1]{4,3,2,0,5,1} fusion(f32[35,1,128,128,128,1]{4,3,2,0,5,1}, f32[1,128,128,128,35]{3,2,1,4,0}), kind=kLoop, calls=%fused_computation, metadata={op_type="add" op_name="jit(fac_conv)/scan/while/body/add" source_file="decomposable_conv.py" source_line=48}
2021-06-25 16:08:31.361137: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.82% %fusion.4 = f32[1,128,128,133,1]{3,2,1,4,0} fusion(f32[35,128,128,128]{3,2,1,0}, s32[]), kind=kLoop, calls=%fused_computation.4
2021-06-25 16:08:31.361150: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.00% %fusion.3 = f32[35,1,1,1,6]{4,3,2,1,0} fusion(f32[35,35,1,1,6]{4,1,3,2,0}, s32[]), kind=kLoop, calls=%fused_computation.3, metadata={op_type="reshape" op_name="jit(fac_conv)/scan/while/body/jit(vmap(reg_conv))/reshape[ dimensions=None\n                                                           new_sizes=(35, 1, 1, 1, 6) ]" source_file="decomposable_conv.py" source_line=38}
2021-06-25 16:08:31.361162: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.00% %fusion.2 = f32[35,1,1,6,1]{4,3,2,1,0} fusion(f32[35,35,1,6,1]{3,1,4,2,0}, s32[]), kind=kLoop, calls=%fused_computation.2, metadata={op_type="reshape" op_name="jit(fac_conv)/scan/while/body/jit(vmap(reg_conv))/reshape[ dimensions=None\n                                                           new_sizes=(35, 1, 1, 6, 1) ]" source_file="decomposable_conv.py" source_line=38}
2021-06-25 16:08:31.361174: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.00% %fusion.1 = f32[35,1,6,1,1]{4,3,2,1,0} fusion(f32[35,35,6,1,1]{2,1,4,3,0}, s32[]), kind=kLoop, calls=%fused_computation.1, metadata={op_type="reshape" op_name="jit(fac_conv)/scan/while/body/jit(vmap(reg_conv))/reshape[ dimensions=None\n                                                           new_sizes=(35, 1, 6, 1, 1) ]" source_file="decomposable_conv.py" source_line=38}
2021-06-25 16:08:31.361186: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]      0 ( 0.00% Σ100.00%)   ... (5 more categories)
2021-06-25 16:08:31.361197: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.361209: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] Execution profile for jit_fac_conv.214: (393 ms @ f_nom)
2021-06-25 16:08:31.361222: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]       607468540 cycles (100.% 100Σ) ::     393183.5 usec (      1430.2 optimal) ::      186.68MFLOP/s ::                    ::      7.05GiB/s ::         4B/cycle :: [total] [entry]
2021-06-25 16:08:31.361237: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]       603757249 cycles (99.39% 99Σ) ::     390781.4 usec                        ::      187.83MFLOP/s ::                    ::      4.99GiB/s ::         3B/cycle :: %while.188 = (f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, s32[], f32[35,1,128,128,128,1]{4,3,2,0,5,1}) while((f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, s32[], f32[35,1,128,128,128,1]{4,3,2,0,5,1}) %tuple.11), condition=%cond_computation.177, body=%body_computation.42, metadata={op_type="while" op_name="jit(fac_conv)/scan/while[ body_nconsts=4\n                          cond_nconsts=0 ]" source_file="decomposable_conv.py" source_line=68}, backend_config="{\"known_trip_count\":{\"n\":\"35\"}}"
2021-06-25 16:08:31.361250: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]         1719669 cycles ( 0.28% 100Σ) ::       1113.1 usec (       953.3 optimal) ::                    ::                    ::    491.33GiB/s ::       341B/cycle :: %fusion.6 = f32[1,128,128,128,35]{4,3,2,1,0} fusion(f32[35,1,128,128,128,1]{4,3,2,0,5,1} %get-tuple-element.194), kind=kLoop, calls=%fused_computation.6, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 2, 3, 4, 0) ]" source_file="decomposable_conv.py" source_line=57}
2021-06-25 16:08:31.361268: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]         1505743 cycles ( 0.25% 100Σ) ::        974.6 usec (       476.6 optimal) ::                    ::                    ::    280.57GiB/s ::       194B/cycle :: %broadcast.21 = f32[35,1,128,128,128,1]{4,3,2,0,5,1} broadcast(f32[] %constant_7), dimensions={}, metadata={op_type="broadcast_in_dim" op_name="jit(fac_conv)/broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4, 5)\n                                shape=(35, 1, 128, 128, 128, 1) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.361281: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]           42220 cycles ( 0.01% 100Σ) ::         27.3 usec (         0.0 optimal) ::                    ::                    ::     285.9KiB/s ::     0.000B/cycle :: %copy.18 = s32[] copy(s32[] %constant_14)
2021-06-25 16:08:31.361294: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]           28131 cycles ( 0.00% 100Σ) ::         18.2 usec (         0.1 optimal) ::                    ::                    ::      3.01GiB/s ::         2B/cycle :: %fusion.9 = f32[35,35,6,1,1]{2,1,4,3,0} fusion(f32[35,35,6,1,1]{4,3,2,1,0} %parameter.3), kind=kLoop, calls=%fused_computation.9, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 0, 2, 3, 4) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.361307: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]           10777 cycles ( 0.00% 100Σ) ::          7.0 usec (         0.0 optimal) ::                    ::                    ::      6.56MiB/s ::     0.004B/cycle :: %tuple.11 = (f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, s32[], f32[35,1,128,128,128,1]{4,3,2,0,5,1}) tuple(f32[35,128,128,128]{3,2,1,0} %parameter.2, f32[35,35,6,1,1]{2,1,4,3,0} %fusion.9, f32[35,35,1,6,1]{3,1,4,2,0} %fusion.8, f32[35,35,1,1,6]{4,1,3,2,0} %fusion.7, s32[] %copy.18, f32[35,1,128,128,128,1]{4,3,2,0,5,1} %broadcast.21)
2021-06-25 16:08:31.361320: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]            8403 cycles ( 0.00% 100Σ) ::          5.4 usec (         0.1 optimal) ::                    ::                    ::     10.07GiB/s ::         6B/cycle :: %fusion.8 = f32[35,35,1,6,1]{3,1,4,2,0} fusion(f32[35,35,1,6,1]{4,3,2,1,0} %parameter.4), kind=kLoop, calls=%fused_computation.8, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 0, 2, 3, 4) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.361333: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]            7068 cycles ( 0.00% 100Σ) ::          4.6 usec (         0.1 optimal) ::                    ::                    ::     11.97GiB/s ::         8B/cycle :: %fusion.7 = f32[35,35,1,1,6]{4,1,3,2,0} fusion(f32[35,35,1,1,6]{4,3,2,1,0} %parameter.5), kind=kLoop, calls=%fused_computation.7, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 0, 2, 3, 4) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.361344: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]               0 cycles ( 0.00% 100Σ) ::          0.0 usec (         0.0 optimal) ::                    ::                    ::                ::                  :: %tuple.213 = (f32[1,128,128,128,35]{4,3,2,1,0}) tuple(f32[1,128,128,128,35]{4,3,2,1,0} %fusion.6)
2021-06-25 16:08:31.361356: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.361368: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] ********** microseconds above estimated optimum report **********
2021-06-25 16:08:31.361380: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 720 microseconds above estimated optimum in total.
2021-06-25 16:08:31.361392: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 0 microseconds above estimated optimum (-0.00%) not accounted for by the data.
2021-06-25 16:08:31.361409: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 7 ops.
2021-06-25 16:08:31.361421: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.361432: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] ********** categories table for microseconds above estimated optimum **********
2021-06-25 16:08:31.361444: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.361455: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]  498 (69.16% Σ69.16%)   broadcast (1 ops)
2021-06-25 16:08:31.361468: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 69.16% %broadcast.21 = f32[35,1,128,128,128,1]{4,3,2,0,5,1} broadcast(f32[]), dimensions={}, metadata={op_type="broadcast_in_dim" op_name="jit(fac_conv)/broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4, 5)\n                                shape=(35, 1, 128, 128, 128, 1) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.361480: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]  188 (26.07% Σ95.24%)   loop fusion (4 ops)
2021-06-25 16:08:31.361492: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 22.19% %fusion.6 = f32[1,128,128,128,35]{4,3,2,1,0} fusion(f32[35,1,128,128,128,1]{4,3,2,0,5,1}), kind=kLoop, calls=%fused_computation.6, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 2, 3, 4, 0) ]" source_file="decomposable_conv.py" source_line=57}
2021-06-25 16:08:31.361505: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  2.52% %fusion.9 = f32[35,35,6,1,1]{2,1,4,3,0} fusion(f32[35,35,6,1,1]{4,3,2,1,0}), kind=kLoop, calls=%fused_computation.9, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 0, 2, 3, 4) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.361517: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.74% %fusion.8 = f32[35,35,1,6,1]{3,1,4,2,0} fusion(f32[35,35,1,6,1]{4,3,2,1,0}), kind=kLoop, calls=%fused_computation.8, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 0, 2, 3, 4) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.361529: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.62% %fusion.7 = f32[35,35,1,1,6]{4,1,3,2,0} fusion(f32[35,35,1,1,6]{4,3,2,1,0}), kind=kLoop, calls=%fused_computation.7, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 0, 2, 3, 4) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.361540: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]   27 ( 3.80% Σ99.03%)   data formatting (1 ops)
2021-06-25 16:08:31.361552: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  3.80% %copy.18 = s32[] copy(s32[])
2021-06-25 16:08:31.361564: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]    7 ( 0.97% Σ100.00%)   tuple (1 ops)
2021-06-25 16:08:31.361576: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.97% %tuple.11 = (f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, s32[], f32[35,1,128,128,128,1]{4,3,2,0,5,1}) tuple(f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, ...(+2))
2021-06-25 16:08:31.361587: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.361603: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.361615: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] ********** microseconds report **********
2021-06-25 16:08:31.361627: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 393,184 microseconds in total.
2021-06-25 16:08:31.361639: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 252 microseconds ( 0.06%) not accounted for by the data.
2021-06-25 16:08:31.361650: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 9 ops.
2021-06-25 16:08:31.361662: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.361674: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] ********** categories table for microseconds **********
2021-06-25 16:08:31.361685: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.361696: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]  390,781 (99.39% Σ99.45%)   while (1 ops)
2021-06-25 16:08:31.361709: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 99.39% %while.188 = (f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, s32[], f32[35,1,128,128,128,1]{4,3,2,0,5,1}) while((f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, s32[], f32[35,1,128,128,128,1]{4,3,2,0,5,1})), condition=%cond_computation.177, body=%body_computation.42, metadata={op_type="while" op_name="jit(fac_conv)/scan/while[ body_nconsts=4\n                          cond_nconsts=0 ]" source_file="decomposable_conv.py" source_line=68}, backend_config="{\"known_trip_count\":{\"n\":\"35\"}}"
2021-06-25 16:08:31.361722: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]    1,141 ( 0.29% Σ99.74%)   loop fusion (4 ops)
2021-06-25 16:08:31.361734: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.28% %fusion.6 = f32[1,128,128,128,35]{4,3,2,1,0} fusion(f32[35,1,128,128,128,1]{4,3,2,0,5,1}), kind=kLoop, calls=%fused_computation.6, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 2, 3, 4, 0) ]" source_file="decomposable_conv.py" source_line=57}
2021-06-25 16:08:31.361746: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.00% %fusion.9 = f32[35,35,6,1,1]{2,1,4,3,0} fusion(f32[35,35,6,1,1]{4,3,2,1,0}), kind=kLoop, calls=%fused_computation.9, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 0, 2, 3, 4) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.361758: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.00% %fusion.8 = f32[35,35,1,6,1]{3,1,4,2,0} fusion(f32[35,35,1,6,1]{4,3,2,1,0}), kind=kLoop, calls=%fused_computation.8, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 0, 2, 3, 4) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.361769: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.00% %fusion.7 = f32[35,35,1,1,6]{4,1,3,2,0} fusion(f32[35,35,1,1,6]{4,3,2,1,0}), kind=kLoop, calls=%fused_computation.7, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 0, 2, 3, 4) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.361781: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]      975 ( 0.25% Σ99.99%)   broadcast (1 ops)
2021-06-25 16:08:31.361793: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.25% %broadcast.21 = f32[35,1,128,128,128,1]{4,3,2,0,5,1} broadcast(f32[]), dimensions={}, metadata={op_type="broadcast_in_dim" op_name="jit(fac_conv)/broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4, 5)\n                                shape=(35, 1, 128, 128, 128, 1) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.361809: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]       27 ( 0.01% Σ100.00%)   data formatting (1 ops)
2021-06-25 16:08:31.361821: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.01% %copy.18 = s32[] copy(s32[])
2021-06-25 16:08:31.361833: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]        7 ( 0.00% Σ100.00%)   tuple (2 ops)
2021-06-25 16:08:31.361845: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.00% %tuple.11 = (f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, s32[], f32[35,1,128,128,128,1]{4,3,2,0,5,1}) tuple(f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, ...(+2))
2021-06-25 16:08:31.361857: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.00% %tuple.213 = (f32[1,128,128,128,35]{4,3,2,1,0}) tuple(f32[1,128,128,128,35]{4,3,2,1,0})
2021-06-25 16:08:31.361869: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.361880: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.361891: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] ********** MiB read+written report **********
2021-06-25 16:08:31.361903: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 2,838 MiB read+written in total.
2021-06-25 16:08:31.361915: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 0 MiB read+written ( 0.00%) not accounted for by the data.
2021-06-25 16:08:31.361927: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] There are 16 ops.
2021-06-25 16:08:31.361938: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.361950: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] ********** categories table for MiB read+written **********
2021-06-25 16:08:31.361961: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276] 
2021-06-25 16:08:31.361973: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]  1,998 (70.40% Σ70.40%)   while (1 ops)
2021-06-25 16:08:31.361985: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 70.40% %while.188 = (f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, s32[], f32[35,1,128,128,128,1]{4,3,2,0,5,1}) while((f32[35,128,128,128]{3,2,1,0}, f32[35,35,6,1,1]{2,1,4,3,0}, f32[35,35,1,6,1]{3,1,4,2,0}, f32[35,35,1,1,6]{4,1,3,2,0}, s32[], f32[35,1,128,128,128,1]{4,3,2,0,5,1})), condition=%cond_computation.177, body=%body_computation.42, metadata={op_type="while" op_name="jit(fac_conv)/scan/while[ body_nconsts=4\n                          cond_nconsts=0 ]" source_file="decomposable_conv.py" source_line=68}, backend_config="{\"known_trip_count\":{\"n\":\"35\"}}"
2021-06-25 16:08:31.361998: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]    560 (19.74% Σ90.14%)   loop fusion (4 ops)
2021-06-25 16:08:31.362010: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               * 19.73% %fusion.6 = f32[1,128,128,128,35]{4,3,2,1,0} fusion(f32[35,1,128,128,128,1]{4,3,2,0,5,1}), kind=kLoop, calls=%fused_computation.6, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 2, 3, 4, 0) ]" source_file="decomposable_conv.py" source_line=57}
2021-06-25 16:08:31.362026: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.00% %fusion.7 = f32[35,35,1,1,6]{4,1,3,2,0} fusion(f32[35,35,1,1,6]{4,3,2,1,0}), kind=kLoop, calls=%fused_computation.7, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 0, 2, 3, 4) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.362038: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.00% %fusion.8 = f32[35,35,1,6,1]{3,1,4,2,0} fusion(f32[35,35,1,6,1]{4,3,2,1,0}), kind=kLoop, calls=%fused_computation.8, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 0, 2, 3, 4) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.362050: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  0.00% %fusion.9 = f32[35,35,6,1,1]{2,1,4,3,0} fusion(f32[35,35,6,1,1]{4,3,2,1,0}), kind=kLoop, calls=%fused_computation.9, metadata={op_type="transpose" op_name="jit(fac_conv)/transpose[ permutation=(1, 0, 2, 3, 4) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.362062: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]    280 ( 9.86% Σ100.00%)   broadcast (1 ops)
2021-06-25 16:08:31.362074: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]                               *  9.86% %broadcast.21 = f32[35,1,128,128,128,1]{4,3,2,0,5,1} broadcast(f32[]), dimensions={}, metadata={op_type="broadcast_in_dim" op_name="jit(fac_conv)/broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4, 5)\n                                shape=(35, 1, 128, 128, 128, 1) ]" source_file="decomposable_conv.py" source_line=51}
2021-06-25 16:08:31.362086: I external/org_tensorflow/tensorflow/compiler/xla/service/executable.cc:276]      0 ( 0.00% Σ100.00%)   ... (5 more categories)