output of jax.jit(mjx.kinematics).lower(mjx_m, mjx_d).as_text()

module @jit_kinematics attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<9x3xf64> {mhlo.layout_mode = "default"}, %arg2: tensor<9x4xf64> {mhlo.layout_mode = "default"}, %arg3: tensor<9x3xf64> {mhlo.layout_mode = "default"}, %arg4: tensor<9x4xf64> {mhlo.layout_mode = "default"}, %arg5: tensor<7x3xf64> {mhlo.layout_mode = "default"}, %arg6: tensor<7x3xf64> {mhlo.layout_mode = "default"}, %arg7: tensor<61x3xf64> {mhlo.layout_mode = "default"}, %arg8: tensor<61x4xf64> {mhlo.layout_mode = "default"}, %arg9: tensor<1x3xf64> {mhlo.layout_mode = "default"}, %arg10: tensor<1x4xf64> {mhlo.layout_mode = "default"}, %arg11: tensor<i64> {mhlo.layout_mode = "default"}, %arg12: tensor<f64> {mhlo.layout_mode = "default"}, %arg13: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg14: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg15: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg16: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg17: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg18: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg19: tensor<9x6xf64> {mhlo.layout_mode = "default"}, %arg20: tensor<0xui8> {mhlo.layout_mode = "default"}, %arg21: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg22: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg23: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg24: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg25: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg26: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg27: tensor<0x3xf64> {mhlo.layout_mode = "default"}, %arg28: tensor<0x3x3xf64> {mhlo.layout_mode = "default"}, %arg29: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg30: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg31: tensor<9x3xf64> {mhlo.layout_mode = "default"}, %arg32: tensor<7x6xf64> {mhlo.layout_mode = "default"}, %arg33: tensor<9x10xf64> {mhlo.layout_mode = "default"}, %arg34: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg35: tensor<0x6xf64> {mhlo.layout_mode = "default"}, %arg36: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg37: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg38: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg39: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg40: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg41: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg42: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg43: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg44: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg45: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg46: tensor<0x7xf64> {mhlo.layout_mode = "default"}, %arg47: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg48: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg49: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg50: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg51: tensor<7x7xf64> {mhlo.layout_mode = "default"}, %arg52: tensor<9x10xf64> {mhlo.layout_mode = "default"}, %arg53: tensor<7x7xf64> {mhlo.layout_mode = "default"}, %arg54: tensor<7x7xf64> {mhlo.layout_mode = "default"}, %arg55: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg56: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg57: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg58: tensor<0xui8> {mhlo.layout_mode = "default"}, %arg59: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg60: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg61: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg62: tensor<9x6xf64> {mhlo.layout_mode = "default"}, %arg63: tensor<7x6xf64> {mhlo.layout_mode = "default"}, %arg64: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg65: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg66: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg67: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg68: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg69: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg70: tensor<9x3xf64> {mhlo.layout_mode = "default"}, %arg71: tensor<9x3xf64> {mhlo.layout_mode = "default"}, %arg72: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg73: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg74: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg75: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg76: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg77: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg78: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg79: tensor<0xi32> {mhlo.layout_mode = "default"}, %arg80: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg81: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg82: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg83: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg84: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg85: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg86: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg87: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg88: tensor<9x6xf64> {mhlo.layout_mode = "default"}, %arg89: tensor<9x6xf64> {mhlo.layout_mode = "default"}, %arg90: tensor<9x6xf64> {mhlo.layout_mode = "default"}, %arg91: tensor<396xf64> {mhlo.layout_mode = "default"}, %arg92: tensor<396x3xf64> {mhlo.layout_mode = "default"}, %arg93: tensor<396x3x3xf64> {mhlo.layout_mode = "default"}, %arg94: tensor<396xf64> {mhlo.layout_mode = "default"}, %arg95: tensor<396x5xf64> {mhlo.layout_mode = "default"}, %arg96: tensor<396x2xf64> {mhlo.layout_mode = "default"}, %arg97: tensor<396x2xf64> {mhlo.layout_mode = "default"}, %arg98: tensor<396x5xf64> {mhlo.layout_mode = "default"}, %arg99: tensor<396xi32> {mhlo.layout_mode = "default"}, %arg100: tensor<396xi32> {mhlo.layout_mode = "default"}, %arg101: tensor<396x2xi32> {mhlo.layout_mode = "default"}, %arg102: tensor<1591xi64> {mhlo.layout_mode = "default"}, %arg103: tensor<1591x7xf64> {mhlo.layout_mode = "default"}, %arg104: tensor<1591xf64> {mhlo.layout_mode = "default"}, %arg105: tensor<1591xf64> {mhlo.layout_mode = "default"}, %arg106: tensor<1591xf64> {mhlo.layout_mode = "default"}, %arg107: tensor<1591xf64> {mhlo.layout_mode = "default"}, %arg108: tensor<1591xf64> {mhlo.layout_mode = "default"}, %arg109: tensor<1591xf64> {mhlo.layout_mode = "default"}, %arg110: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg111: tensor<0xf64> {mhlo.layout_mode = "default"}, %arg112: tensor<0xf64> {mhlo.layout_mode = "default"}) -> (tensor<i64> {jax.result_info = ".solver_niter", mhlo.layout_mode = "default"}, tensor<f64> {jax.result_info = ".time", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qpos", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qvel", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".act", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qacc_warmstart", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".ctrl", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qfrc_applied", mhlo.layout_mode = "default"}, tensor<9x6xf64> {jax.result_info = ".xfrc_applied", mhlo.layout_mode = "default"}, tensor<0xui8> {jax.result_info = ".eq_active", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".mocap_pos", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".mocap_quat", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qacc", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".act_dot", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".userdata", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".sensordata", mhlo.layout_mode = "default"}, tensor<9x3xf64> {jax.result_info = ".xpos", mhlo.layout_mode = "default"}, tensor<9x4xf64> {jax.result_info = ".xquat", mhlo.layout_mode = "default"}, tensor<9x3x3xf64> {jax.result_info = ".xmat", mhlo.layout_mode = "default"}, tensor<9x3xf64> {jax.result_info = ".xipos", mhlo.layout_mode = "default"}, tensor<9x3x3xf64> {jax.result_info = ".ximat", mhlo.layout_mode = "default"}, tensor<7x3xf64> {jax.result_info = ".xanchor", mhlo.layout_mode = "default"}, tensor<7x3xf64> {jax.result_info = ".xaxis", mhlo.layout_mode = "default"}, tensor<61x3xf64> {jax.result_info = ".geom_xpos", mhlo.layout_mode = "default"}, tensor<61x3x3xf64> {jax.result_info = ".geom_xmat", mhlo.layout_mode = "default"}, tensor<1x3xf64> {jax.result_info = ".site_xpos", mhlo.layout_mode = "default"}, tensor<1x3x3xf64> {jax.result_info = ".site_xmat", mhlo.layout_mode = "default"}, tensor<0x3xf64> {jax.result_info = ".cam_xpos", mhlo.layout_mode = "default"}, tensor<0x3x3xf64> {jax.result_info = ".cam_xmat", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".light_xpos", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".light_xdir", mhlo.layout_mode = "default"}, tensor<9x3xf64> {jax.result_info = ".subtree_com", mhlo.layout_mode = "default"}, tensor<7x6xf64> {jax.result_info = ".cdof", mhlo.layout_mode = "default"}, tensor<9x10xf64> {jax.result_info = ".cinert", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".flexvert_xpos", mhlo.layout_mode = "default"}, tensor<0x6xf64> {jax.result_info = ".flexelem_aabb", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".flexedge_J_rownnz", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".flexedge_J_rowadr", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".flexedge_J_colind", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".flexedge_J", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".flexedge_length", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".ten_wrapadr", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".ten_wrapnum", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".ten_J_rownnz", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".ten_J_rowadr", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".ten_J_colind", mhlo.layout_mode = "default"}, tensor<0x7xf64> {jax.result_info = ".ten_J", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".ten_length", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".wrap_obj", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".wrap_xpos", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".actuator_length", mhlo.layout_mode = "default"}, tensor<7x7xf64> {jax.result_info = ".actuator_moment", mhlo.layout_mode = "default"}, tensor<9x10xf64> {jax.result_info = ".crb", mhlo.layout_mode = "default"}, tensor<7x7xf64> {jax.result_info = ".qM", mhlo.layout_mode = "default"}, tensor<7x7xf64> {jax.result_info = ".qLD", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".qLDiagInv", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qLDiagSqrtInv", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".bvh_aabb_dyn", mhlo.layout_mode = "default"}, tensor<0xui8> {jax.result_info = ".bvh_active", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".flexedge_velocity", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".ten_velocity", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".actuator_velocity", mhlo.layout_mode = "default"}, tensor<9x6xf64> {jax.result_info = ".cvel", mhlo.layout_mode = "default"}, tensor<7x6xf64> {jax.result_info = ".cdof_dot", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qfrc_bias", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".qfrc_spring", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".qfrc_damper", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qfrc_gravcomp", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qfrc_fluid", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qfrc_passive", mhlo.layout_mode = "default"}, tensor<9x3xf64> {jax.result_info = ".subtree_linvel", mhlo.layout_mode = "default"}, tensor<9x3xf64> {jax.result_info = ".subtree_angmom", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".qH", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".qHDiagInv", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".D_rownnz", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".D_rowadr", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".D_colind", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".B_rownnz", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".B_rowadr", mhlo.layout_mode = "default"}, tensor<0xi32> {jax.result_info = ".B_colind", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".qDeriv", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = ".qLU", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qfrc_actuator", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".actuator_force", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qfrc_smooth", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qacc_smooth", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qfrc_constraint", mhlo.layout_mode = "default"}, tensor<7xf64> {jax.result_info = ".qfrc_inverse", mhlo.layout_mode = "default"}, tensor<9x6xf64> {jax.result_info = ".cacc", mhlo.layout_mode = "default"}, tensor<9x6xf64> {jax.result_info = ".cfrc_int", mhlo.layout_mode = "default"}, tensor<9x6xf64> {jax.result_info = ".cfrc_ext", mhlo.layout_mode = "default"}, tensor<396xf64> {jax.result_info = ".contact.dist", mhlo.layout_mode = "default"}, tensor<396x3xf64> {jax.result_info = ".contact.pos", mhlo.layout_mode = "default"}, tensor<396x3x3xf64> {jax.result_info = ".contact.frame", mhlo.layout_mode = "default"}, tensor<396xf64> {jax.result_info = ".contact.includemargin", mhlo.layout_mode = "default"}, tensor<396x5xf64> {jax.result_info = ".contact.friction", mhlo.layout_mode = "default"}, tensor<396x2xf64> {jax.result_info = ".contact.solref", mhlo.layout_mode = "default"}, tensor<396x2xf64> {jax.result_info = ".contact.solreffriction", mhlo.layout_mode = "default"}, tensor<396x5xf64> {jax.result_info = ".contact.solimp", mhlo.layout_mode = "default"}, tensor<396xi32> {jax.result_info = ".contact.geom1", mhlo.layout_mode = "default"}, tensor<396xi32> {jax.result_info = ".contact.geom2", mhlo.layout_mode = "default"}, tensor<396x2xi32> {jax.result_info = ".contact.geom", mhlo.layout_mode = "default"}, tensor<1591xi64> {jax.result_info = ".efc_type", mhlo.layout_mode = "default"}, tensor<1591x7xf64> {jax.result_info = ".efc_J", mhlo.layout_mode = "default"}, tensor<1591xf64> {jax.result_info = ".efc_pos", mhlo.layout_mode = "default"}, tensor<1591xf64> {jax.result_info = ".efc_margin", mhlo.layout_mode = "default"}, tensor<1591xf64> {jax.result_info = ".efc_frictionloss", mhlo.layout_mode = "default"}, tensor<1591xf64> {jax.result_info = ".efc_D", mhlo.layout_mode = "default"}, tensor<1591xf64> {jax.result_info = ".efc_aref", mhlo.layout_mode = "default"}, tensor<1591xf64> {jax.result_info = ".efc_force", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = "._qM_sparse", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = "._qLD_sparse", mhlo.layout_mode = "default"}, tensor<0xf64> {jax.result_info = "._qLDiagInv_sparse", mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<0> : tensor<1xi64>
    %c_0 = stablehlo.constant dense<0> : tensor<1xi64>
    %c_1 = stablehlo.constant dense<0> : tensor<1xi64>
    %c_2 = stablehlo.constant dense<0> : tensor<1x1xi64>
    %c_3 = stablehlo.constant dense<0> : tensor<1x1xi64>
    %c_4 = stablehlo.constant dense<0> : tensor<1x1xi64>
    %c_5 = stablehlo.constant dense<0> : tensor<1x1xi64>
    %c_6 = stablehlo.constant dense<0> : tensor<1xi64>
    %c_7 = stablehlo.constant dense<1> : tensor<1x1xi64>
    %c_8 = stablehlo.constant dense<1> : tensor<1x1xi64>
    %c_9 = stablehlo.constant dense<1> : tensor<1x1xi64>
    %c_10 = stablehlo.constant dense<1> : tensor<1x1xi64>
    %c_11 = stablehlo.constant dense<0> : tensor<1xi64>
    %c_12 = stablehlo.constant dense<2> : tensor<1x1xi64>
    %c_13 = stablehlo.constant dense<2> : tensor<1x1xi64>
    %c_14 = stablehlo.constant dense<2> : tensor<1x1xi64>
    %c_15 = stablehlo.constant dense<2> : tensor<1x1xi64>
    %c_16 = stablehlo.constant dense<0> : tensor<1xi64>
    %c_17 = stablehlo.constant dense<3> : tensor<1x1xi64>
    %c_18 = stablehlo.constant dense<3> : tensor<1x1xi64>
    %c_19 = stablehlo.constant dense<3> : tensor<1x1xi64>
    %c_20 = stablehlo.constant dense<3> : tensor<1x1xi64>
    %c_21 = stablehlo.constant dense<0> : tensor<1xi64>
    %c_22 = stablehlo.constant dense<4> : tensor<1x1xi64>
    %c_23 = stablehlo.constant dense<4> : tensor<1x1xi64>
    %c_24 = stablehlo.constant dense<4> : tensor<1x1xi64>
    %c_25 = stablehlo.constant dense<4> : tensor<1x1xi64>
    %c_26 = stablehlo.constant dense<0> : tensor<1xi64>
    %c_27 = stablehlo.constant dense<5> : tensor<1x1xi64>
    %c_28 = stablehlo.constant dense<5> : tensor<1x1xi64>
    %c_29 = stablehlo.constant dense<5> : tensor<1x1xi64>
    %c_30 = stablehlo.constant dense<5> : tensor<1x1xi64>
    %c_31 = stablehlo.constant dense<0> : tensor<1xi64>
    %c_32 = stablehlo.constant dense<6> : tensor<1x1xi64>
    %c_33 = stablehlo.constant dense<6> : tensor<1x1xi64>
    %c_34 = stablehlo.constant dense<6> : tensor<1x1xi64>
    %c_35 = stablehlo.constant dense<6> : tensor<1x1xi64>
    %c_36 = stablehlo.constant dense<[0, 1, 2, 3, 4, 5, 6]> : tensor<7xi64>
    %c_37 = stablehlo.constant dense<[0, 1, 2, 3, 4, 5, 6]> : tensor<7xi64>
    %c_38 = stablehlo.constant dense<[0, 1, 2, 3, 4, 5, 6]> : tensor<7xi64>
    %c_39 = stablehlo.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8]> : tensor<9xi64>
    %c_40 = stablehlo.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8]> : tensor<9xi64>
    %c_41 = stablehlo.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8]> : tensor<9xi64>
    %c_42 = stablehlo.constant dense<[1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 8, 8]> : tensor<61xi32>
    %c_43 = stablehlo.constant dense<8> : tensor<1xi32>
    %0 = call @_take() : () -> tensor<1x0xf64>
    %1 = call @_take_0(%arg1, %c) : (tensor<9x3xf64>, tensor<1xi64>) -> tensor<1x3xf64>
    %2 = call @_take_1(%arg2, %c) : (tensor<9x4xf64>, tensor<1xi64>) -> tensor<1x4xf64>
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %3 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f64>) -> tensor<0x3xf64>
    %cst_44 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %4 = stablehlo.broadcast_in_dim %cst_44, dims = [] : (tensor<f64>) -> tensor<0x3xf64>
    %5 = stablehlo.broadcast_in_dim %2, dims = [0, 1] : (tensor<1x4xf64>) -> tensor<1x4x1xf64>
    %6 = stablehlo.broadcast_in_dim %2, dims = [0, 2] : (tensor<1x4xf64>) -> tensor<1x1x4xf64>
    %7 = stablehlo.broadcast_in_dim %5, dims = [0, 1, 2] : (tensor<1x4x1xf64>) -> tensor<1x4x4xf64>
    %8 = stablehlo.broadcast_in_dim %6, dims = [0, 1, 2] : (tensor<1x1x4xf64>) -> tensor<1x4x4xf64>
    %9 = stablehlo.multiply %7, %8 : tensor<1x4x4xf64>
    %10 = stablehlo.slice %9 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %11 = stablehlo.reshape %10 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %12 = stablehlo.slice %9 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %13 = stablehlo.reshape %12 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %14 = stablehlo.add %11, %13 : tensor<1xf64>
    %15 = stablehlo.slice %9 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %16 = stablehlo.reshape %15 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %17 = stablehlo.subtract %14, %16 : tensor<1xf64>
    %18 = stablehlo.slice %9 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %19 = stablehlo.reshape %18 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %20 = stablehlo.subtract %17, %19 : tensor<1xf64>
    %21 = stablehlo.slice %9 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %22 = stablehlo.reshape %21 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %23 = stablehlo.slice %9 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %24 = stablehlo.reshape %23 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %25 = stablehlo.subtract %22, %24 : tensor<1xf64>
    %cst_45 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %26 = stablehlo.broadcast_in_dim %cst_45, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %27 = stablehlo.multiply %26, %25 : tensor<1xf64>
    %28 = stablehlo.slice %9 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %29 = stablehlo.reshape %28 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %30 = stablehlo.slice %9 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %31 = stablehlo.reshape %30 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %32 = stablehlo.add %29, %31 : tensor<1xf64>
    %cst_46 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %33 = stablehlo.broadcast_in_dim %cst_46, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %34 = stablehlo.multiply %33, %32 : tensor<1xf64>
    %35 = stablehlo.slice %9 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %36 = stablehlo.reshape %35 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %37 = stablehlo.slice %9 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %38 = stablehlo.reshape %37 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %39 = stablehlo.add %36, %38 : tensor<1xf64>
    %cst_47 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %40 = stablehlo.broadcast_in_dim %cst_47, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %41 = stablehlo.multiply %40, %39 : tensor<1xf64>
    %42 = stablehlo.slice %9 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %43 = stablehlo.reshape %42 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %44 = stablehlo.slice %9 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %45 = stablehlo.reshape %44 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %46 = stablehlo.subtract %43, %45 : tensor<1xf64>
    %47 = stablehlo.slice %9 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %48 = stablehlo.reshape %47 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %49 = stablehlo.add %46, %48 : tensor<1xf64>
    %50 = stablehlo.slice %9 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %51 = stablehlo.reshape %50 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %52 = stablehlo.subtract %49, %51 : tensor<1xf64>
    %53 = stablehlo.slice %9 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %54 = stablehlo.reshape %53 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %55 = stablehlo.slice %9 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %56 = stablehlo.reshape %55 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %57 = stablehlo.subtract %54, %56 : tensor<1xf64>
    %cst_48 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %58 = stablehlo.broadcast_in_dim %cst_48, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %59 = stablehlo.multiply %58, %57 : tensor<1xf64>
    %60 = stablehlo.slice %9 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %61 = stablehlo.reshape %60 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %62 = stablehlo.slice %9 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %63 = stablehlo.reshape %62 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %64 = stablehlo.subtract %61, %63 : tensor<1xf64>
    %cst_49 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %65 = stablehlo.broadcast_in_dim %cst_49, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %66 = stablehlo.multiply %65, %64 : tensor<1xf64>
    %67 = stablehlo.slice %9 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %68 = stablehlo.reshape %67 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %69 = stablehlo.slice %9 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %70 = stablehlo.reshape %69 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %71 = stablehlo.add %68, %70 : tensor<1xf64>
    %cst_50 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %72 = stablehlo.broadcast_in_dim %cst_50, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %73 = stablehlo.multiply %72, %71 : tensor<1xf64>
    %74 = stablehlo.slice %9 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %75 = stablehlo.reshape %74 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %76 = stablehlo.slice %9 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %77 = stablehlo.reshape %76 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %78 = stablehlo.subtract %75, %77 : tensor<1xf64>
    %79 = stablehlo.slice %9 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %80 = stablehlo.reshape %79 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %81 = stablehlo.subtract %78, %80 : tensor<1xf64>
    %82 = stablehlo.slice %9 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %83 = stablehlo.reshape %82 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %84 = stablehlo.add %81, %83 : tensor<1xf64>
    %85 = stablehlo.broadcast_in_dim %20, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %86 = stablehlo.broadcast_in_dim %27, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %87 = stablehlo.broadcast_in_dim %34, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %88 = stablehlo.concatenate %85, %86, %87, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %89 = stablehlo.broadcast_in_dim %41, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %90 = stablehlo.broadcast_in_dim %52, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %91 = stablehlo.broadcast_in_dim %59, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %92 = stablehlo.concatenate %89, %90, %91, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %93 = stablehlo.broadcast_in_dim %66, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %94 = stablehlo.broadcast_in_dim %73, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %95 = stablehlo.broadcast_in_dim %84, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %96 = stablehlo.concatenate %93, %94, %95, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %97 = stablehlo.broadcast_in_dim %88, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %98 = stablehlo.broadcast_in_dim %92, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %99 = stablehlo.broadcast_in_dim %96, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %100 = stablehlo.concatenate %97, %98, %99, dim = 1 : (tensor<1x1x3xf64>, tensor<1x1x3xf64>, tensor<1x1x3xf64>) -> tensor<1x3x3xf64>
    %101 = stablehlo.broadcast_in_dim %3, dims = [1, 2] : (tensor<0x3xf64>) -> tensor<1x0x3xf64>
    %102 = stablehlo.broadcast_in_dim %4, dims = [1, 2] : (tensor<0x3xf64>) -> tensor<1x0x3xf64>
    %103 = call @_take_2(%1, %c_0) : (tensor<1x3xf64>, tensor<1xi64>) -> tensor<1x3xf64>
    %104 = call @_take_3(%2, %c_0) : (tensor<1x4xf64>, tensor<1xi64>) -> tensor<1x4xf64>
    %105 = call @_take() : () -> tensor<1x0xf64>
    %106 = stablehlo.slice %arg1 [1:2, 0:3] : (tensor<9x3xf64>) -> tensor<1x3xf64>
    %107 = stablehlo.slice %arg2 [1:2, 0:4] : (tensor<9x4xf64>) -> tensor<1x4xf64>
    %108 = stablehlo.slice %104 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %109 = stablehlo.reshape %108 : (tensor<1x1xf64>) -> tensor<1xf64>
    %110 = stablehlo.slice %104 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %111 = stablehlo.dot_general %110, %106, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %112 = stablehlo.broadcast_in_dim %111, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %113 = stablehlo.broadcast_in_dim %112, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %114 = stablehlo.multiply %113, %110 : tensor<1x3xf64>
    %cst_51 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %115 = stablehlo.broadcast_in_dim %cst_51, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %116 = stablehlo.multiply %115, %114 : tensor<1x3xf64>
    %117 = stablehlo.multiply %109, %109 : tensor<1xf64>
    %118 = stablehlo.dot_general %110, %110, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %119 = stablehlo.subtract %117, %118 : tensor<1xf64>
    %120 = stablehlo.broadcast_in_dim %119, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %121 = stablehlo.broadcast_in_dim %120, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %122 = stablehlo.multiply %121, %106 : tensor<1x3xf64>
    %123 = stablehlo.add %116, %122 : tensor<1x3xf64>
    %cst_52 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %124 = stablehlo.broadcast_in_dim %cst_52, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %125 = stablehlo.multiply %124, %109 : tensor<1xf64>
    %126 = call @cross(%110, %106) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %127 = stablehlo.broadcast_in_dim %125, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %128 = stablehlo.broadcast_in_dim %127, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %129 = stablehlo.multiply %128, %126 : tensor<1x3xf64>
    %130 = stablehlo.add %123, %129 : tensor<1x3xf64>
    %131 = stablehlo.add %103, %130 : tensor<1x3xf64>
    %132 = stablehlo.slice %104 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %133 = stablehlo.reshape %132 : (tensor<1x1xf64>) -> tensor<1xf64>
    %134 = stablehlo.slice %107 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %135 = stablehlo.reshape %134 : (tensor<1x1xf64>) -> tensor<1xf64>
    %136 = stablehlo.multiply %133, %135 : tensor<1xf64>
    %137 = stablehlo.slice %104 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %138 = stablehlo.reshape %137 : (tensor<1x1xf64>) -> tensor<1xf64>
    %139 = stablehlo.slice %107 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %140 = stablehlo.reshape %139 : (tensor<1x1xf64>) -> tensor<1xf64>
    %141 = stablehlo.multiply %138, %140 : tensor<1xf64>
    %142 = stablehlo.subtract %136, %141 : tensor<1xf64>
    %143 = stablehlo.slice %104 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %144 = stablehlo.reshape %143 : (tensor<1x1xf64>) -> tensor<1xf64>
    %145 = stablehlo.slice %107 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %146 = stablehlo.reshape %145 : (tensor<1x1xf64>) -> tensor<1xf64>
    %147 = stablehlo.multiply %144, %146 : tensor<1xf64>
    %148 = stablehlo.subtract %142, %147 : tensor<1xf64>
    %149 = stablehlo.slice %104 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %150 = stablehlo.reshape %149 : (tensor<1x1xf64>) -> tensor<1xf64>
    %151 = stablehlo.slice %107 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %152 = stablehlo.reshape %151 : (tensor<1x1xf64>) -> tensor<1xf64>
    %153 = stablehlo.multiply %150, %152 : tensor<1xf64>
    %154 = stablehlo.subtract %148, %153 : tensor<1xf64>
    %155 = stablehlo.slice %104 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %156 = stablehlo.reshape %155 : (tensor<1x1xf64>) -> tensor<1xf64>
    %157 = stablehlo.slice %107 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %158 = stablehlo.reshape %157 : (tensor<1x1xf64>) -> tensor<1xf64>
    %159 = stablehlo.multiply %156, %158 : tensor<1xf64>
    %160 = stablehlo.slice %104 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %161 = stablehlo.reshape %160 : (tensor<1x1xf64>) -> tensor<1xf64>
    %162 = stablehlo.slice %107 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %163 = stablehlo.reshape %162 : (tensor<1x1xf64>) -> tensor<1xf64>
    %164 = stablehlo.multiply %161, %163 : tensor<1xf64>
    %165 = stablehlo.add %159, %164 : tensor<1xf64>
    %166 = stablehlo.slice %104 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %167 = stablehlo.reshape %166 : (tensor<1x1xf64>) -> tensor<1xf64>
    %168 = stablehlo.slice %107 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %169 = stablehlo.reshape %168 : (tensor<1x1xf64>) -> tensor<1xf64>
    %170 = stablehlo.multiply %167, %169 : tensor<1xf64>
    %171 = stablehlo.add %165, %170 : tensor<1xf64>
    %172 = stablehlo.slice %104 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %173 = stablehlo.reshape %172 : (tensor<1x1xf64>) -> tensor<1xf64>
    %174 = stablehlo.slice %107 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %175 = stablehlo.reshape %174 : (tensor<1x1xf64>) -> tensor<1xf64>
    %176 = stablehlo.multiply %173, %175 : tensor<1xf64>
    %177 = stablehlo.subtract %171, %176 : tensor<1xf64>
    %178 = stablehlo.slice %104 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %179 = stablehlo.reshape %178 : (tensor<1x1xf64>) -> tensor<1xf64>
    %180 = stablehlo.slice %107 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %181 = stablehlo.reshape %180 : (tensor<1x1xf64>) -> tensor<1xf64>
    %182 = stablehlo.multiply %179, %181 : tensor<1xf64>
    %183 = stablehlo.slice %104 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %184 = stablehlo.reshape %183 : (tensor<1x1xf64>) -> tensor<1xf64>
    %185 = stablehlo.slice %107 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %186 = stablehlo.reshape %185 : (tensor<1x1xf64>) -> tensor<1xf64>
    %187 = stablehlo.multiply %184, %186 : tensor<1xf64>
    %188 = stablehlo.subtract %182, %187 : tensor<1xf64>
    %189 = stablehlo.slice %104 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %190 = stablehlo.reshape %189 : (tensor<1x1xf64>) -> tensor<1xf64>
    %191 = stablehlo.slice %107 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %192 = stablehlo.reshape %191 : (tensor<1x1xf64>) -> tensor<1xf64>
    %193 = stablehlo.multiply %190, %192 : tensor<1xf64>
    %194 = stablehlo.add %188, %193 : tensor<1xf64>
    %195 = stablehlo.slice %104 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %196 = stablehlo.reshape %195 : (tensor<1x1xf64>) -> tensor<1xf64>
    %197 = stablehlo.slice %107 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %198 = stablehlo.reshape %197 : (tensor<1x1xf64>) -> tensor<1xf64>
    %199 = stablehlo.multiply %196, %198 : tensor<1xf64>
    %200 = stablehlo.add %194, %199 : tensor<1xf64>
    %201 = stablehlo.slice %104 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %202 = stablehlo.reshape %201 : (tensor<1x1xf64>) -> tensor<1xf64>
    %203 = stablehlo.slice %107 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %204 = stablehlo.reshape %203 : (tensor<1x1xf64>) -> tensor<1xf64>
    %205 = stablehlo.multiply %202, %204 : tensor<1xf64>
    %206 = stablehlo.slice %104 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %207 = stablehlo.reshape %206 : (tensor<1x1xf64>) -> tensor<1xf64>
    %208 = stablehlo.slice %107 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %209 = stablehlo.reshape %208 : (tensor<1x1xf64>) -> tensor<1xf64>
    %210 = stablehlo.multiply %207, %209 : tensor<1xf64>
    %211 = stablehlo.add %205, %210 : tensor<1xf64>
    %212 = stablehlo.slice %104 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %213 = stablehlo.reshape %212 : (tensor<1x1xf64>) -> tensor<1xf64>
    %214 = stablehlo.slice %107 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %215 = stablehlo.reshape %214 : (tensor<1x1xf64>) -> tensor<1xf64>
    %216 = stablehlo.multiply %213, %215 : tensor<1xf64>
    %217 = stablehlo.subtract %211, %216 : tensor<1xf64>
    %218 = stablehlo.slice %104 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %219 = stablehlo.reshape %218 : (tensor<1x1xf64>) -> tensor<1xf64>
    %220 = stablehlo.slice %107 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %221 = stablehlo.reshape %220 : (tensor<1x1xf64>) -> tensor<1xf64>
    %222 = stablehlo.multiply %219, %221 : tensor<1xf64>
    %223 = stablehlo.add %217, %222 : tensor<1xf64>
    %224 = stablehlo.broadcast_in_dim %154, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %225 = stablehlo.broadcast_in_dim %177, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %226 = stablehlo.broadcast_in_dim %200, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %227 = stablehlo.broadcast_in_dim %223, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %228 = stablehlo.concatenate %224, %225, %226, %227, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %cst_53 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %229 = stablehlo.broadcast_in_dim %cst_53, dims = [] : (tensor<f64>) -> tensor<0x3xf64>
    %cst_54 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %230 = stablehlo.broadcast_in_dim %cst_54, dims = [] : (tensor<f64>) -> tensor<0x3xf64>
    %231 = stablehlo.broadcast_in_dim %228, dims = [0, 1] : (tensor<1x4xf64>) -> tensor<1x4x1xf64>
    %232 = stablehlo.broadcast_in_dim %228, dims = [0, 2] : (tensor<1x4xf64>) -> tensor<1x1x4xf64>
    %233 = stablehlo.broadcast_in_dim %231, dims = [0, 1, 2] : (tensor<1x4x1xf64>) -> tensor<1x4x4xf64>
    %234 = stablehlo.broadcast_in_dim %232, dims = [0, 1, 2] : (tensor<1x1x4xf64>) -> tensor<1x4x4xf64>
    %235 = stablehlo.multiply %233, %234 : tensor<1x4x4xf64>
    %236 = stablehlo.slice %235 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %237 = stablehlo.reshape %236 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %238 = stablehlo.slice %235 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %239 = stablehlo.reshape %238 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %240 = stablehlo.add %237, %239 : tensor<1xf64>
    %241 = stablehlo.slice %235 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %242 = stablehlo.reshape %241 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %243 = stablehlo.subtract %240, %242 : tensor<1xf64>
    %244 = stablehlo.slice %235 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %245 = stablehlo.reshape %244 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %246 = stablehlo.subtract %243, %245 : tensor<1xf64>
    %247 = stablehlo.slice %235 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %248 = stablehlo.reshape %247 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %249 = stablehlo.slice %235 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %250 = stablehlo.reshape %249 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %251 = stablehlo.subtract %248, %250 : tensor<1xf64>
    %cst_55 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %252 = stablehlo.broadcast_in_dim %cst_55, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %253 = stablehlo.multiply %252, %251 : tensor<1xf64>
    %254 = stablehlo.slice %235 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %255 = stablehlo.reshape %254 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %256 = stablehlo.slice %235 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %257 = stablehlo.reshape %256 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %258 = stablehlo.add %255, %257 : tensor<1xf64>
    %cst_56 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %259 = stablehlo.broadcast_in_dim %cst_56, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %260 = stablehlo.multiply %259, %258 : tensor<1xf64>
    %261 = stablehlo.slice %235 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %262 = stablehlo.reshape %261 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %263 = stablehlo.slice %235 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %264 = stablehlo.reshape %263 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %265 = stablehlo.add %262, %264 : tensor<1xf64>
    %cst_57 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %266 = stablehlo.broadcast_in_dim %cst_57, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %267 = stablehlo.multiply %266, %265 : tensor<1xf64>
    %268 = stablehlo.slice %235 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %269 = stablehlo.reshape %268 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %270 = stablehlo.slice %235 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %271 = stablehlo.reshape %270 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %272 = stablehlo.subtract %269, %271 : tensor<1xf64>
    %273 = stablehlo.slice %235 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %274 = stablehlo.reshape %273 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %275 = stablehlo.add %272, %274 : tensor<1xf64>
    %276 = stablehlo.slice %235 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %277 = stablehlo.reshape %276 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %278 = stablehlo.subtract %275, %277 : tensor<1xf64>
    %279 = stablehlo.slice %235 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %280 = stablehlo.reshape %279 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %281 = stablehlo.slice %235 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %282 = stablehlo.reshape %281 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %283 = stablehlo.subtract %280, %282 : tensor<1xf64>
    %cst_58 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %284 = stablehlo.broadcast_in_dim %cst_58, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %285 = stablehlo.multiply %284, %283 : tensor<1xf64>
    %286 = stablehlo.slice %235 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %287 = stablehlo.reshape %286 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %288 = stablehlo.slice %235 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %289 = stablehlo.reshape %288 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %290 = stablehlo.subtract %287, %289 : tensor<1xf64>
    %cst_59 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %291 = stablehlo.broadcast_in_dim %cst_59, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %292 = stablehlo.multiply %291, %290 : tensor<1xf64>
    %293 = stablehlo.slice %235 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %294 = stablehlo.reshape %293 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %295 = stablehlo.slice %235 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %296 = stablehlo.reshape %295 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %297 = stablehlo.add %294, %296 : tensor<1xf64>
    %cst_60 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %298 = stablehlo.broadcast_in_dim %cst_60, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %299 = stablehlo.multiply %298, %297 : tensor<1xf64>
    %300 = stablehlo.slice %235 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %301 = stablehlo.reshape %300 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %302 = stablehlo.slice %235 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %303 = stablehlo.reshape %302 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %304 = stablehlo.subtract %301, %303 : tensor<1xf64>
    %305 = stablehlo.slice %235 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %306 = stablehlo.reshape %305 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %307 = stablehlo.subtract %304, %306 : tensor<1xf64>
    %308 = stablehlo.slice %235 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %309 = stablehlo.reshape %308 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %310 = stablehlo.add %307, %309 : tensor<1xf64>
    %311 = stablehlo.broadcast_in_dim %246, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %312 = stablehlo.broadcast_in_dim %253, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %313 = stablehlo.broadcast_in_dim %260, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %314 = stablehlo.concatenate %311, %312, %313, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %315 = stablehlo.broadcast_in_dim %267, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %316 = stablehlo.broadcast_in_dim %278, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %317 = stablehlo.broadcast_in_dim %285, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %318 = stablehlo.concatenate %315, %316, %317, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %319 = stablehlo.broadcast_in_dim %292, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %320 = stablehlo.broadcast_in_dim %299, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %321 = stablehlo.broadcast_in_dim %310, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %322 = stablehlo.concatenate %319, %320, %321, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %323 = stablehlo.broadcast_in_dim %314, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %324 = stablehlo.broadcast_in_dim %318, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %325 = stablehlo.broadcast_in_dim %322, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %326 = stablehlo.concatenate %323, %324, %325, dim = 1 : (tensor<1x1x3xf64>, tensor<1x1x3xf64>, tensor<1x1x3xf64>) -> tensor<1x3x3xf64>
    %327 = stablehlo.broadcast_in_dim %229, dims = [1, 2] : (tensor<0x3xf64>) -> tensor<1x0x3xf64>
    %328 = stablehlo.broadcast_in_dim %230, dims = [1, 2] : (tensor<0x3xf64>) -> tensor<1x0x3xf64>
    %329 = call @_take_2(%131, %c_1) : (tensor<1x3xf64>, tensor<1xi64>) -> tensor<1x3xf64>
    %330 = call @_take_3(%228, %c_1) : (tensor<1x4xf64>, tensor<1xi64>) -> tensor<1x4xf64>
    %331 = call @_take_4(%arg5, %c_2) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %332 = call @_take_4(%arg6, %c_3) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %333 = call @_take_6(%arg13, %c_4) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %334 = call @_take_6(%arg0, %c_5) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %335 = stablehlo.slice %arg1 [2:3, 0:3] : (tensor<9x3xf64>) -> tensor<1x3xf64>
    %336 = stablehlo.slice %arg2 [2:3, 0:4] : (tensor<9x4xf64>) -> tensor<1x4xf64>
    %337 = stablehlo.slice %330 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %338 = stablehlo.reshape %337 : (tensor<1x1xf64>) -> tensor<1xf64>
    %339 = stablehlo.slice %330 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %340 = stablehlo.dot_general %339, %335, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %341 = stablehlo.broadcast_in_dim %340, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %342 = stablehlo.broadcast_in_dim %341, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %343 = stablehlo.multiply %342, %339 : tensor<1x3xf64>
    %cst_61 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %344 = stablehlo.broadcast_in_dim %cst_61, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %345 = stablehlo.multiply %344, %343 : tensor<1x3xf64>
    %346 = stablehlo.multiply %338, %338 : tensor<1xf64>
    %347 = stablehlo.dot_general %339, %339, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %348 = stablehlo.subtract %346, %347 : tensor<1xf64>
    %349 = stablehlo.broadcast_in_dim %348, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %350 = stablehlo.broadcast_in_dim %349, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %351 = stablehlo.multiply %350, %335 : tensor<1x3xf64>
    %352 = stablehlo.add %345, %351 : tensor<1x3xf64>
    %cst_62 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %353 = stablehlo.broadcast_in_dim %cst_62, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %354 = stablehlo.multiply %353, %338 : tensor<1xf64>
    %355 = call @cross(%339, %335) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %356 = stablehlo.broadcast_in_dim %354, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %357 = stablehlo.broadcast_in_dim %356, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %358 = stablehlo.multiply %357, %355 : tensor<1x3xf64>
    %359 = stablehlo.add %352, %358 : tensor<1x3xf64>
    %360 = stablehlo.add %329, %359 : tensor<1x3xf64>
    %361 = stablehlo.slice %330 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %362 = stablehlo.reshape %361 : (tensor<1x1xf64>) -> tensor<1xf64>
    %363 = stablehlo.slice %336 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %364 = stablehlo.reshape %363 : (tensor<1x1xf64>) -> tensor<1xf64>
    %365 = stablehlo.multiply %362, %364 : tensor<1xf64>
    %366 = stablehlo.slice %330 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %367 = stablehlo.reshape %366 : (tensor<1x1xf64>) -> tensor<1xf64>
    %368 = stablehlo.slice %336 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %369 = stablehlo.reshape %368 : (tensor<1x1xf64>) -> tensor<1xf64>
    %370 = stablehlo.multiply %367, %369 : tensor<1xf64>
    %371 = stablehlo.subtract %365, %370 : tensor<1xf64>
    %372 = stablehlo.slice %330 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %373 = stablehlo.reshape %372 : (tensor<1x1xf64>) -> tensor<1xf64>
    %374 = stablehlo.slice %336 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %375 = stablehlo.reshape %374 : (tensor<1x1xf64>) -> tensor<1xf64>
    %376 = stablehlo.multiply %373, %375 : tensor<1xf64>
    %377 = stablehlo.subtract %371, %376 : tensor<1xf64>
    %378 = stablehlo.slice %330 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %379 = stablehlo.reshape %378 : (tensor<1x1xf64>) -> tensor<1xf64>
    %380 = stablehlo.slice %336 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %381 = stablehlo.reshape %380 : (tensor<1x1xf64>) -> tensor<1xf64>
    %382 = stablehlo.multiply %379, %381 : tensor<1xf64>
    %383 = stablehlo.subtract %377, %382 : tensor<1xf64>
    %384 = stablehlo.slice %330 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %385 = stablehlo.reshape %384 : (tensor<1x1xf64>) -> tensor<1xf64>
    %386 = stablehlo.slice %336 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %387 = stablehlo.reshape %386 : (tensor<1x1xf64>) -> tensor<1xf64>
    %388 = stablehlo.multiply %385, %387 : tensor<1xf64>
    %389 = stablehlo.slice %330 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %390 = stablehlo.reshape %389 : (tensor<1x1xf64>) -> tensor<1xf64>
    %391 = stablehlo.slice %336 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %392 = stablehlo.reshape %391 : (tensor<1x1xf64>) -> tensor<1xf64>
    %393 = stablehlo.multiply %390, %392 : tensor<1xf64>
    %394 = stablehlo.add %388, %393 : tensor<1xf64>
    %395 = stablehlo.slice %330 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %396 = stablehlo.reshape %395 : (tensor<1x1xf64>) -> tensor<1xf64>
    %397 = stablehlo.slice %336 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %398 = stablehlo.reshape %397 : (tensor<1x1xf64>) -> tensor<1xf64>
    %399 = stablehlo.multiply %396, %398 : tensor<1xf64>
    %400 = stablehlo.add %394, %399 : tensor<1xf64>
    %401 = stablehlo.slice %330 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %402 = stablehlo.reshape %401 : (tensor<1x1xf64>) -> tensor<1xf64>
    %403 = stablehlo.slice %336 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %404 = stablehlo.reshape %403 : (tensor<1x1xf64>) -> tensor<1xf64>
    %405 = stablehlo.multiply %402, %404 : tensor<1xf64>
    %406 = stablehlo.subtract %400, %405 : tensor<1xf64>
    %407 = stablehlo.slice %330 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %408 = stablehlo.reshape %407 : (tensor<1x1xf64>) -> tensor<1xf64>
    %409 = stablehlo.slice %336 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %410 = stablehlo.reshape %409 : (tensor<1x1xf64>) -> tensor<1xf64>
    %411 = stablehlo.multiply %408, %410 : tensor<1xf64>
    %412 = stablehlo.slice %330 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %413 = stablehlo.reshape %412 : (tensor<1x1xf64>) -> tensor<1xf64>
    %414 = stablehlo.slice %336 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %415 = stablehlo.reshape %414 : (tensor<1x1xf64>) -> tensor<1xf64>
    %416 = stablehlo.multiply %413, %415 : tensor<1xf64>
    %417 = stablehlo.subtract %411, %416 : tensor<1xf64>
    %418 = stablehlo.slice %330 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %419 = stablehlo.reshape %418 : (tensor<1x1xf64>) -> tensor<1xf64>
    %420 = stablehlo.slice %336 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %421 = stablehlo.reshape %420 : (tensor<1x1xf64>) -> tensor<1xf64>
    %422 = stablehlo.multiply %419, %421 : tensor<1xf64>
    %423 = stablehlo.add %417, %422 : tensor<1xf64>
    %424 = stablehlo.slice %330 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %425 = stablehlo.reshape %424 : (tensor<1x1xf64>) -> tensor<1xf64>
    %426 = stablehlo.slice %336 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %427 = stablehlo.reshape %426 : (tensor<1x1xf64>) -> tensor<1xf64>
    %428 = stablehlo.multiply %425, %427 : tensor<1xf64>
    %429 = stablehlo.add %423, %428 : tensor<1xf64>
    %430 = stablehlo.slice %330 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %431 = stablehlo.reshape %430 : (tensor<1x1xf64>) -> tensor<1xf64>
    %432 = stablehlo.slice %336 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %433 = stablehlo.reshape %432 : (tensor<1x1xf64>) -> tensor<1xf64>
    %434 = stablehlo.multiply %431, %433 : tensor<1xf64>
    %435 = stablehlo.slice %330 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %436 = stablehlo.reshape %435 : (tensor<1x1xf64>) -> tensor<1xf64>
    %437 = stablehlo.slice %336 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %438 = stablehlo.reshape %437 : (tensor<1x1xf64>) -> tensor<1xf64>
    %439 = stablehlo.multiply %436, %438 : tensor<1xf64>
    %440 = stablehlo.add %434, %439 : tensor<1xf64>
    %441 = stablehlo.slice %330 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %442 = stablehlo.reshape %441 : (tensor<1x1xf64>) -> tensor<1xf64>
    %443 = stablehlo.slice %336 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %444 = stablehlo.reshape %443 : (tensor<1x1xf64>) -> tensor<1xf64>
    %445 = stablehlo.multiply %442, %444 : tensor<1xf64>
    %446 = stablehlo.subtract %440, %445 : tensor<1xf64>
    %447 = stablehlo.slice %330 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %448 = stablehlo.reshape %447 : (tensor<1x1xf64>) -> tensor<1xf64>
    %449 = stablehlo.slice %336 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %450 = stablehlo.reshape %449 : (tensor<1x1xf64>) -> tensor<1xf64>
    %451 = stablehlo.multiply %448, %450 : tensor<1xf64>
    %452 = stablehlo.add %446, %451 : tensor<1xf64>
    %453 = stablehlo.broadcast_in_dim %383, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %454 = stablehlo.broadcast_in_dim %406, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %455 = stablehlo.broadcast_in_dim %429, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %456 = stablehlo.broadcast_in_dim %452, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %457 = stablehlo.concatenate %453, %454, %455, %456, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %458 = stablehlo.slice %331 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %459 = stablehlo.reshape %458 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %460 = stablehlo.slice %457 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %461 = stablehlo.reshape %460 : (tensor<1x1xf64>) -> tensor<1xf64>
    %462 = stablehlo.slice %457 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %463 = stablehlo.dot_general %462, %459, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %464 = stablehlo.broadcast_in_dim %463, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %465 = stablehlo.broadcast_in_dim %464, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %466 = stablehlo.multiply %465, %462 : tensor<1x3xf64>
    %cst_63 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %467 = stablehlo.broadcast_in_dim %cst_63, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %468 = stablehlo.multiply %467, %466 : tensor<1x3xf64>
    %469 = stablehlo.multiply %461, %461 : tensor<1xf64>
    %470 = stablehlo.dot_general %462, %462, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %471 = stablehlo.subtract %469, %470 : tensor<1xf64>
    %472 = stablehlo.broadcast_in_dim %471, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %473 = stablehlo.broadcast_in_dim %472, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %474 = stablehlo.multiply %473, %459 : tensor<1x3xf64>
    %475 = stablehlo.add %468, %474 : tensor<1x3xf64>
    %cst_64 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %476 = stablehlo.broadcast_in_dim %cst_64, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %477 = stablehlo.multiply %476, %461 : tensor<1xf64>
    %478 = call @cross(%462, %459) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %479 = stablehlo.broadcast_in_dim %477, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %480 = stablehlo.broadcast_in_dim %479, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %481 = stablehlo.multiply %480, %478 : tensor<1x3xf64>
    %482 = stablehlo.add %475, %481 : tensor<1x3xf64>
    %483 = stablehlo.add %482, %360 : tensor<1x3xf64>
    %484 = stablehlo.slice %332 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %485 = stablehlo.reshape %484 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %486 = stablehlo.slice %457 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %487 = stablehlo.reshape %486 : (tensor<1x1xf64>) -> tensor<1xf64>
    %488 = stablehlo.slice %457 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %489 = stablehlo.dot_general %488, %485, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %490 = stablehlo.broadcast_in_dim %489, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %491 = stablehlo.broadcast_in_dim %490, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %492 = stablehlo.multiply %491, %488 : tensor<1x3xf64>
    %cst_65 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %493 = stablehlo.broadcast_in_dim %cst_65, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %494 = stablehlo.multiply %493, %492 : tensor<1x3xf64>
    %495 = stablehlo.multiply %487, %487 : tensor<1xf64>
    %496 = stablehlo.dot_general %488, %488, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %497 = stablehlo.subtract %495, %496 : tensor<1xf64>
    %498 = stablehlo.broadcast_in_dim %497, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %499 = stablehlo.broadcast_in_dim %498, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %500 = stablehlo.multiply %499, %485 : tensor<1x3xf64>
    %501 = stablehlo.add %494, %500 : tensor<1x3xf64>
    %cst_66 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %502 = stablehlo.broadcast_in_dim %cst_66, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %503 = stablehlo.multiply %502, %487 : tensor<1xf64>
    %504 = call @cross(%488, %485) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %505 = stablehlo.broadcast_in_dim %503, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %506 = stablehlo.broadcast_in_dim %505, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %507 = stablehlo.multiply %506, %504 : tensor<1x3xf64>
    %508 = stablehlo.add %501, %507 : tensor<1x3xf64>
    %509 = stablehlo.slice %333 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %510 = stablehlo.reshape %509 : (tensor<1x1xf64>) -> tensor<1xf64>
    %511 = stablehlo.slice %334 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %512 = stablehlo.reshape %511 : (tensor<1x1xf64>) -> tensor<1xf64>
    %513 = stablehlo.subtract %510, %512 : tensor<1xf64>
    %514 = stablehlo.slice %332 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %515 = stablehlo.reshape %514 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %cst_67 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %516 = stablehlo.broadcast_in_dim %cst_67, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %517 = stablehlo.multiply %513, %516 : tensor<1xf64>
    %518 = stablehlo.sine %517 : tensor<1xf64>
    %cst_68 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %519 = stablehlo.broadcast_in_dim %cst_68, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %520 = stablehlo.multiply %513, %519 : tensor<1xf64>
    %521 = stablehlo.cosine %520 : tensor<1xf64>
    %522 = stablehlo.broadcast_in_dim %518, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %523 = stablehlo.broadcast_in_dim %522, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %524 = stablehlo.multiply %515, %523 : tensor<1x3xf64>
    %525 = stablehlo.broadcast_in_dim %521, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %c_69 = stablehlo.constant dense<0> : tensor<i64>
    %526 = stablehlo.reshape %c_69 : (tensor<i64>) -> tensor<1xi64>
    %527 = stablehlo.slice %526 [0:1] : (tensor<1xi64>) -> tensor<1xi64>
    %528 = stablehlo.reshape %527 : (tensor<1xi64>) -> tensor<i64>
    %529 = stablehlo.broadcast_in_dim %528, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %cst_70 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %530 = stablehlo.broadcast_in_dim %cst_70, dims = [] : (tensor<f64>) -> tensor<4xf64>
    %c_71 = stablehlo.constant dense<0> : tensor<i64>
    %531 = stablehlo.broadcast_in_dim %c_71, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %532 = stablehlo.compare  LT, %529, %531,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_72 = stablehlo.constant dense<3> : tensor<i64>
    %533 = stablehlo.broadcast_in_dim %c_72, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %534 = stablehlo.add %529, %533 : tensor<1xi64>
    %535 = call @_where_7(%532, %534, %529) : (tensor<1xi1>, tensor<1xi64>, tensor<1xi64>) -> tensor<1xi64>
    %c_73 = stablehlo.constant dense<0> : tensor<i64>
    %c_74 = stablehlo.constant dense<3> : tensor<i64>
    %536 = call @clip(%535, %c_73, %c_74) : (tensor<1xi64>, tensor<i64>, tensor<i64>) -> tensor<1xi64>
    %537 = call @argsort(%536) : (tensor<1xi64>) -> tensor<1xi64>
    %538 = stablehlo.iota dim = 0 : tensor<1xi64>
    %c_75 = stablehlo.constant dense<0> : tensor<i64>
    %539 = stablehlo.broadcast_in_dim %c_75, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %540 = stablehlo.compare  LT, %537, %539,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_76 = stablehlo.constant dense<1> : tensor<i64>
    %541 = stablehlo.broadcast_in_dim %c_76, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %542 = stablehlo.add %537, %541 : tensor<1xi64>
    %543 = stablehlo.select %540, %542, %537 : tensor<1xi1>, tensor<1xi64>
    %544 = stablehlo.convert %543 : (tensor<1xi64>) -> tensor<1xi32>
    %545 = stablehlo.broadcast_in_dim %544, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %546 = stablehlo.convert %536 : tensor<1xi64>
    %547 = "stablehlo.scatter"(%546, %545, %538) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<1xi64>, tensor<1x1xi32>, tensor<1xi64>) -> tensor<1xi64>
    %548 = stablehlo.convert %547 : tensor<1xi64>
    %c_77 = stablehlo.constant dense<true> : tensor<i1>
    %549 = stablehlo.broadcast_in_dim %c_77, dims = [] : (tensor<i1>) -> tensor<4xi1>
    %c_78 = stablehlo.constant dense<0> : tensor<i64>
    %550 = stablehlo.broadcast_in_dim %c_78, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %551 = stablehlo.compare  LT, %548, %550,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_79 = stablehlo.constant dense<4> : tensor<i64>
    %552 = stablehlo.broadcast_in_dim %c_79, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %553 = stablehlo.add %548, %552 : tensor<1xi64>
    %554 = stablehlo.select %551, %553, %548 : tensor<1xi1>, tensor<1xi64>
    %555 = stablehlo.convert %554 : (tensor<1xi64>) -> tensor<1xi32>
    %556 = stablehlo.broadcast_in_dim %555, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %c_80 = stablehlo.constant dense<false> : tensor<i1>
    %557 = stablehlo.broadcast_in_dim %c_80, dims = [] : (tensor<i1>) -> tensor<1xi1>
    %558 = "stablehlo.scatter"(%549, %556, %557) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i1>, %arg114: tensor<i1>):
      stablehlo.return %arg114 : tensor<i1>
    }) : (tensor<4xi1>, tensor<1x1xi32>, tensor<1xi1>) -> tensor<4xi1>
    %559 = call @_cumulative_reduction(%558) : (tensor<4xi1>) -> tensor<4xi64>
    %c_81 = stablehlo.constant dense<0> : tensor<i64>
    %560 = stablehlo.broadcast_in_dim %c_81, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %c_82 = stablehlo.constant dense<0> : tensor<i64>
    %561 = call @clip_8(%559, %c_82) : (tensor<4xi64>, tensor<i64>) -> tensor<4xi64>
    %c_83 = stablehlo.constant dense<1> : tensor<i64>
    %c_84 = stablehlo.constant dense<0> : tensor<i64>
    %562 = stablehlo.broadcast_in_dim %c_84, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %563 = stablehlo.compare  LT, %561, %562,  SIGNED : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1>
    %c_85 = stablehlo.constant dense<3> : tensor<i64>
    %564 = stablehlo.broadcast_in_dim %c_85, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %565 = stablehlo.add %561, %564 : tensor<4xi64>
    %566 = stablehlo.select %563, %565, %561 : tensor<4xi1>, tensor<4xi64>
    %567 = stablehlo.convert %566 : (tensor<4xi64>) -> tensor<4xi32>
    %568 = stablehlo.broadcast_in_dim %567, dims = [0] : (tensor<4xi32>) -> tensor<4x1xi32>
    %569 = stablehlo.broadcast_in_dim %c_83, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %570 = "stablehlo.scatter"(%560, %568, %569) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<3xi64>, tensor<4x1xi32>, tensor<4xi64>) -> tensor<3xi64>
    %571 = call @_cumulative_reduction_9(%570) : (tensor<3xi64>) -> tensor<3xi64>
    %c_86 = stablehlo.constant dense<1> : tensor<i64>
    %572 = call @floor_divide(%571, %c_86) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_87 = stablehlo.constant dense<4> : tensor<i64>
    %573 = call @remainder_12(%572, %c_87) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_88 = stablehlo.constant dense<0> : tensor<i64>
    %574 = stablehlo.broadcast_in_dim %c_88, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %575 = stablehlo.compare  LT, %548, %574,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_89 = stablehlo.constant dense<4> : tensor<i64>
    %576 = stablehlo.broadcast_in_dim %c_89, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %577 = stablehlo.add %548, %576 : tensor<1xi64>
    %578 = stablehlo.select %575, %577, %548 : tensor<1xi1>, tensor<1xi64>
    %579 = stablehlo.convert %578 : (tensor<1xi64>) -> tensor<1xi32>
    %580 = stablehlo.broadcast_in_dim %579, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %581 = stablehlo.broadcast_in_dim %530, dims = [1] : (tensor<4xf64>) -> tensor<1x4xf64>
    %582 = "stablehlo.scatter"(%581, %580, %525) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<1x1xi32>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %c_90 = stablehlo.constant dense<0> : tensor<i64>
    %583 = stablehlo.broadcast_in_dim %c_90, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %584 = stablehlo.compare  LT, %573, %583,  SIGNED : (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi1>
    %c_91 = stablehlo.constant dense<4> : tensor<i64>
    %585 = stablehlo.broadcast_in_dim %c_91, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %586 = stablehlo.add %573, %585 : tensor<3xi64>
    %587 = stablehlo.select %584, %586, %573 : tensor<3xi1>, tensor<3xi64>
    %588 = stablehlo.convert %587 : (tensor<3xi64>) -> tensor<3xi32>
    %589 = stablehlo.broadcast_in_dim %588, dims = [0] : (tensor<3xi32>) -> tensor<3x1xi32>
    %590 = "stablehlo.scatter"(%582, %589, %524) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<3x1xi32>, tensor<1x3xf64>) -> tensor<1x4xf64>
    %591 = stablehlo.slice %457 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %592 = stablehlo.reshape %591 : (tensor<1x1xf64>) -> tensor<1xf64>
    %593 = stablehlo.slice %590 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %594 = stablehlo.reshape %593 : (tensor<1x1xf64>) -> tensor<1xf64>
    %595 = stablehlo.multiply %592, %594 : tensor<1xf64>
    %596 = stablehlo.slice %457 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %597 = stablehlo.reshape %596 : (tensor<1x1xf64>) -> tensor<1xf64>
    %598 = stablehlo.slice %590 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %599 = stablehlo.reshape %598 : (tensor<1x1xf64>) -> tensor<1xf64>
    %600 = stablehlo.multiply %597, %599 : tensor<1xf64>
    %601 = stablehlo.subtract %595, %600 : tensor<1xf64>
    %602 = stablehlo.slice %457 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %603 = stablehlo.reshape %602 : (tensor<1x1xf64>) -> tensor<1xf64>
    %604 = stablehlo.slice %590 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %605 = stablehlo.reshape %604 : (tensor<1x1xf64>) -> tensor<1xf64>
    %606 = stablehlo.multiply %603, %605 : tensor<1xf64>
    %607 = stablehlo.subtract %601, %606 : tensor<1xf64>
    %608 = stablehlo.slice %457 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %609 = stablehlo.reshape %608 : (tensor<1x1xf64>) -> tensor<1xf64>
    %610 = stablehlo.slice %590 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %611 = stablehlo.reshape %610 : (tensor<1x1xf64>) -> tensor<1xf64>
    %612 = stablehlo.multiply %609, %611 : tensor<1xf64>
    %613 = stablehlo.subtract %607, %612 : tensor<1xf64>
    %614 = stablehlo.slice %457 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %615 = stablehlo.reshape %614 : (tensor<1x1xf64>) -> tensor<1xf64>
    %616 = stablehlo.slice %590 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %617 = stablehlo.reshape %616 : (tensor<1x1xf64>) -> tensor<1xf64>
    %618 = stablehlo.multiply %615, %617 : tensor<1xf64>
    %619 = stablehlo.slice %457 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %620 = stablehlo.reshape %619 : (tensor<1x1xf64>) -> tensor<1xf64>
    %621 = stablehlo.slice %590 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %622 = stablehlo.reshape %621 : (tensor<1x1xf64>) -> tensor<1xf64>
    %623 = stablehlo.multiply %620, %622 : tensor<1xf64>
    %624 = stablehlo.add %618, %623 : tensor<1xf64>
    %625 = stablehlo.slice %457 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %626 = stablehlo.reshape %625 : (tensor<1x1xf64>) -> tensor<1xf64>
    %627 = stablehlo.slice %590 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %628 = stablehlo.reshape %627 : (tensor<1x1xf64>) -> tensor<1xf64>
    %629 = stablehlo.multiply %626, %628 : tensor<1xf64>
    %630 = stablehlo.add %624, %629 : tensor<1xf64>
    %631 = stablehlo.slice %457 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %632 = stablehlo.reshape %631 : (tensor<1x1xf64>) -> tensor<1xf64>
    %633 = stablehlo.slice %590 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %634 = stablehlo.reshape %633 : (tensor<1x1xf64>) -> tensor<1xf64>
    %635 = stablehlo.multiply %632, %634 : tensor<1xf64>
    %636 = stablehlo.subtract %630, %635 : tensor<1xf64>
    %637 = stablehlo.slice %457 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %638 = stablehlo.reshape %637 : (tensor<1x1xf64>) -> tensor<1xf64>
    %639 = stablehlo.slice %590 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %640 = stablehlo.reshape %639 : (tensor<1x1xf64>) -> tensor<1xf64>
    %641 = stablehlo.multiply %638, %640 : tensor<1xf64>
    %642 = stablehlo.slice %457 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %643 = stablehlo.reshape %642 : (tensor<1x1xf64>) -> tensor<1xf64>
    %644 = stablehlo.slice %590 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %645 = stablehlo.reshape %644 : (tensor<1x1xf64>) -> tensor<1xf64>
    %646 = stablehlo.multiply %643, %645 : tensor<1xf64>
    %647 = stablehlo.subtract %641, %646 : tensor<1xf64>
    %648 = stablehlo.slice %457 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %649 = stablehlo.reshape %648 : (tensor<1x1xf64>) -> tensor<1xf64>
    %650 = stablehlo.slice %590 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %651 = stablehlo.reshape %650 : (tensor<1x1xf64>) -> tensor<1xf64>
    %652 = stablehlo.multiply %649, %651 : tensor<1xf64>
    %653 = stablehlo.add %647, %652 : tensor<1xf64>
    %654 = stablehlo.slice %457 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %655 = stablehlo.reshape %654 : (tensor<1x1xf64>) -> tensor<1xf64>
    %656 = stablehlo.slice %590 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %657 = stablehlo.reshape %656 : (tensor<1x1xf64>) -> tensor<1xf64>
    %658 = stablehlo.multiply %655, %657 : tensor<1xf64>
    %659 = stablehlo.add %653, %658 : tensor<1xf64>
    %660 = stablehlo.slice %457 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %661 = stablehlo.reshape %660 : (tensor<1x1xf64>) -> tensor<1xf64>
    %662 = stablehlo.slice %590 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %663 = stablehlo.reshape %662 : (tensor<1x1xf64>) -> tensor<1xf64>
    %664 = stablehlo.multiply %661, %663 : tensor<1xf64>
    %665 = stablehlo.slice %457 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %666 = stablehlo.reshape %665 : (tensor<1x1xf64>) -> tensor<1xf64>
    %667 = stablehlo.slice %590 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %668 = stablehlo.reshape %667 : (tensor<1x1xf64>) -> tensor<1xf64>
    %669 = stablehlo.multiply %666, %668 : tensor<1xf64>
    %670 = stablehlo.add %664, %669 : tensor<1xf64>
    %671 = stablehlo.slice %457 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %672 = stablehlo.reshape %671 : (tensor<1x1xf64>) -> tensor<1xf64>
    %673 = stablehlo.slice %590 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %674 = stablehlo.reshape %673 : (tensor<1x1xf64>) -> tensor<1xf64>
    %675 = stablehlo.multiply %672, %674 : tensor<1xf64>
    %676 = stablehlo.subtract %670, %675 : tensor<1xf64>
    %677 = stablehlo.slice %457 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %678 = stablehlo.reshape %677 : (tensor<1x1xf64>) -> tensor<1xf64>
    %679 = stablehlo.slice %590 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %680 = stablehlo.reshape %679 : (tensor<1x1xf64>) -> tensor<1xf64>
    %681 = stablehlo.multiply %678, %680 : tensor<1xf64>
    %682 = stablehlo.add %676, %681 : tensor<1xf64>
    %683 = stablehlo.broadcast_in_dim %613, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %684 = stablehlo.broadcast_in_dim %636, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %685 = stablehlo.broadcast_in_dim %659, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %686 = stablehlo.broadcast_in_dim %682, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %687 = stablehlo.concatenate %683, %684, %685, %686, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %688 = stablehlo.slice %331 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %689 = stablehlo.reshape %688 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %690 = stablehlo.slice %687 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %691 = stablehlo.reshape %690 : (tensor<1x1xf64>) -> tensor<1xf64>
    %692 = stablehlo.slice %687 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %693 = stablehlo.dot_general %692, %689, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %694 = stablehlo.broadcast_in_dim %693, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %695 = stablehlo.broadcast_in_dim %694, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %696 = stablehlo.multiply %695, %692 : tensor<1x3xf64>
    %cst_92 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %697 = stablehlo.broadcast_in_dim %cst_92, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %698 = stablehlo.multiply %697, %696 : tensor<1x3xf64>
    %699 = stablehlo.multiply %691, %691 : tensor<1xf64>
    %700 = stablehlo.dot_general %692, %692, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %701 = stablehlo.subtract %699, %700 : tensor<1xf64>
    %702 = stablehlo.broadcast_in_dim %701, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %703 = stablehlo.broadcast_in_dim %702, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %704 = stablehlo.multiply %703, %689 : tensor<1x3xf64>
    %705 = stablehlo.add %698, %704 : tensor<1x3xf64>
    %cst_93 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %706 = stablehlo.broadcast_in_dim %cst_93, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %707 = stablehlo.multiply %706, %691 : tensor<1xf64>
    %708 = call @cross(%692, %689) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %709 = stablehlo.broadcast_in_dim %707, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %710 = stablehlo.broadcast_in_dim %709, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %711 = stablehlo.multiply %710, %708 : tensor<1x3xf64>
    %712 = stablehlo.add %705, %711 : tensor<1x3xf64>
    %713 = stablehlo.subtract %483, %712 : tensor<1x3xf64>
    %714 = stablehlo.broadcast_in_dim %483, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %715 = stablehlo.broadcast_in_dim %508, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %716 = stablehlo.broadcast_in_dim %687, dims = [0, 1] : (tensor<1x4xf64>) -> tensor<1x4x1xf64>
    %717 = stablehlo.broadcast_in_dim %687, dims = [0, 2] : (tensor<1x4xf64>) -> tensor<1x1x4xf64>
    %718 = stablehlo.broadcast_in_dim %716, dims = [0, 1, 2] : (tensor<1x4x1xf64>) -> tensor<1x4x4xf64>
    %719 = stablehlo.broadcast_in_dim %717, dims = [0, 1, 2] : (tensor<1x1x4xf64>) -> tensor<1x4x4xf64>
    %720 = stablehlo.multiply %718, %719 : tensor<1x4x4xf64>
    %721 = stablehlo.slice %720 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %722 = stablehlo.reshape %721 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %723 = stablehlo.slice %720 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %724 = stablehlo.reshape %723 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %725 = stablehlo.add %722, %724 : tensor<1xf64>
    %726 = stablehlo.slice %720 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %727 = stablehlo.reshape %726 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %728 = stablehlo.subtract %725, %727 : tensor<1xf64>
    %729 = stablehlo.slice %720 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %730 = stablehlo.reshape %729 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %731 = stablehlo.subtract %728, %730 : tensor<1xf64>
    %732 = stablehlo.slice %720 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %733 = stablehlo.reshape %732 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %734 = stablehlo.slice %720 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %735 = stablehlo.reshape %734 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %736 = stablehlo.subtract %733, %735 : tensor<1xf64>
    %cst_94 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %737 = stablehlo.broadcast_in_dim %cst_94, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %738 = stablehlo.multiply %737, %736 : tensor<1xf64>
    %739 = stablehlo.slice %720 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %740 = stablehlo.reshape %739 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %741 = stablehlo.slice %720 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %742 = stablehlo.reshape %741 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %743 = stablehlo.add %740, %742 : tensor<1xf64>
    %cst_95 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %744 = stablehlo.broadcast_in_dim %cst_95, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %745 = stablehlo.multiply %744, %743 : tensor<1xf64>
    %746 = stablehlo.slice %720 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %747 = stablehlo.reshape %746 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %748 = stablehlo.slice %720 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %749 = stablehlo.reshape %748 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %750 = stablehlo.add %747, %749 : tensor<1xf64>
    %cst_96 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %751 = stablehlo.broadcast_in_dim %cst_96, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %752 = stablehlo.multiply %751, %750 : tensor<1xf64>
    %753 = stablehlo.slice %720 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %754 = stablehlo.reshape %753 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %755 = stablehlo.slice %720 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %756 = stablehlo.reshape %755 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %757 = stablehlo.subtract %754, %756 : tensor<1xf64>
    %758 = stablehlo.slice %720 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %759 = stablehlo.reshape %758 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %760 = stablehlo.add %757, %759 : tensor<1xf64>
    %761 = stablehlo.slice %720 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %762 = stablehlo.reshape %761 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %763 = stablehlo.subtract %760, %762 : tensor<1xf64>
    %764 = stablehlo.slice %720 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %765 = stablehlo.reshape %764 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %766 = stablehlo.slice %720 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %767 = stablehlo.reshape %766 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %768 = stablehlo.subtract %765, %767 : tensor<1xf64>
    %cst_97 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %769 = stablehlo.broadcast_in_dim %cst_97, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %770 = stablehlo.multiply %769, %768 : tensor<1xf64>
    %771 = stablehlo.slice %720 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %772 = stablehlo.reshape %771 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %773 = stablehlo.slice %720 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %774 = stablehlo.reshape %773 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %775 = stablehlo.subtract %772, %774 : tensor<1xf64>
    %cst_98 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %776 = stablehlo.broadcast_in_dim %cst_98, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %777 = stablehlo.multiply %776, %775 : tensor<1xf64>
    %778 = stablehlo.slice %720 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %779 = stablehlo.reshape %778 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %780 = stablehlo.slice %720 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %781 = stablehlo.reshape %780 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %782 = stablehlo.add %779, %781 : tensor<1xf64>
    %cst_99 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %783 = stablehlo.broadcast_in_dim %cst_99, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %784 = stablehlo.multiply %783, %782 : tensor<1xf64>
    %785 = stablehlo.slice %720 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %786 = stablehlo.reshape %785 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %787 = stablehlo.slice %720 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %788 = stablehlo.reshape %787 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %789 = stablehlo.subtract %786, %788 : tensor<1xf64>
    %790 = stablehlo.slice %720 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %791 = stablehlo.reshape %790 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %792 = stablehlo.subtract %789, %791 : tensor<1xf64>
    %793 = stablehlo.slice %720 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %794 = stablehlo.reshape %793 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %795 = stablehlo.add %792, %794 : tensor<1xf64>
    %796 = stablehlo.broadcast_in_dim %731, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %797 = stablehlo.broadcast_in_dim %738, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %798 = stablehlo.broadcast_in_dim %745, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %799 = stablehlo.concatenate %796, %797, %798, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %800 = stablehlo.broadcast_in_dim %752, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %801 = stablehlo.broadcast_in_dim %763, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %802 = stablehlo.broadcast_in_dim %770, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %803 = stablehlo.concatenate %800, %801, %802, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %804 = stablehlo.broadcast_in_dim %777, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %805 = stablehlo.broadcast_in_dim %784, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %806 = stablehlo.broadcast_in_dim %795, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %807 = stablehlo.concatenate %804, %805, %806, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %808 = stablehlo.broadcast_in_dim %799, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %809 = stablehlo.broadcast_in_dim %803, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %810 = stablehlo.broadcast_in_dim %807, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %811 = stablehlo.concatenate %808, %809, %810, dim = 1 : (tensor<1x1x3xf64>, tensor<1x1x3xf64>, tensor<1x1x3xf64>) -> tensor<1x3x3xf64>
    %812 = call @_take_2(%713, %c_6) : (tensor<1x3xf64>, tensor<1xi64>) -> tensor<1x3xf64>
    %813 = call @_take_3(%687, %c_6) : (tensor<1x4xf64>, tensor<1xi64>) -> tensor<1x4xf64>
    %814 = call @_take_4(%arg5, %c_7) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %815 = call @_take_4(%arg6, %c_8) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %816 = call @_take_6(%arg13, %c_9) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %817 = call @_take_6(%arg0, %c_10) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %818 = stablehlo.slice %arg1 [3:4, 0:3] : (tensor<9x3xf64>) -> tensor<1x3xf64>
    %819 = stablehlo.slice %arg2 [3:4, 0:4] : (tensor<9x4xf64>) -> tensor<1x4xf64>
    %820 = stablehlo.slice %813 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %821 = stablehlo.reshape %820 : (tensor<1x1xf64>) -> tensor<1xf64>
    %822 = stablehlo.slice %813 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %823 = stablehlo.dot_general %822, %818, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %824 = stablehlo.broadcast_in_dim %823, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %825 = stablehlo.broadcast_in_dim %824, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %826 = stablehlo.multiply %825, %822 : tensor<1x3xf64>
    %cst_100 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %827 = stablehlo.broadcast_in_dim %cst_100, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %828 = stablehlo.multiply %827, %826 : tensor<1x3xf64>
    %829 = stablehlo.multiply %821, %821 : tensor<1xf64>
    %830 = stablehlo.dot_general %822, %822, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %831 = stablehlo.subtract %829, %830 : tensor<1xf64>
    %832 = stablehlo.broadcast_in_dim %831, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %833 = stablehlo.broadcast_in_dim %832, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %834 = stablehlo.multiply %833, %818 : tensor<1x3xf64>
    %835 = stablehlo.add %828, %834 : tensor<1x3xf64>
    %cst_101 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %836 = stablehlo.broadcast_in_dim %cst_101, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %837 = stablehlo.multiply %836, %821 : tensor<1xf64>
    %838 = call @cross(%822, %818) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %839 = stablehlo.broadcast_in_dim %837, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %840 = stablehlo.broadcast_in_dim %839, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %841 = stablehlo.multiply %840, %838 : tensor<1x3xf64>
    %842 = stablehlo.add %835, %841 : tensor<1x3xf64>
    %843 = stablehlo.add %812, %842 : tensor<1x3xf64>
    %844 = stablehlo.slice %813 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %845 = stablehlo.reshape %844 : (tensor<1x1xf64>) -> tensor<1xf64>
    %846 = stablehlo.slice %819 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %847 = stablehlo.reshape %846 : (tensor<1x1xf64>) -> tensor<1xf64>
    %848 = stablehlo.multiply %845, %847 : tensor<1xf64>
    %849 = stablehlo.slice %813 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %850 = stablehlo.reshape %849 : (tensor<1x1xf64>) -> tensor<1xf64>
    %851 = stablehlo.slice %819 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %852 = stablehlo.reshape %851 : (tensor<1x1xf64>) -> tensor<1xf64>
    %853 = stablehlo.multiply %850, %852 : tensor<1xf64>
    %854 = stablehlo.subtract %848, %853 : tensor<1xf64>
    %855 = stablehlo.slice %813 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %856 = stablehlo.reshape %855 : (tensor<1x1xf64>) -> tensor<1xf64>
    %857 = stablehlo.slice %819 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %858 = stablehlo.reshape %857 : (tensor<1x1xf64>) -> tensor<1xf64>
    %859 = stablehlo.multiply %856, %858 : tensor<1xf64>
    %860 = stablehlo.subtract %854, %859 : tensor<1xf64>
    %861 = stablehlo.slice %813 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %862 = stablehlo.reshape %861 : (tensor<1x1xf64>) -> tensor<1xf64>
    %863 = stablehlo.slice %819 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %864 = stablehlo.reshape %863 : (tensor<1x1xf64>) -> tensor<1xf64>
    %865 = stablehlo.multiply %862, %864 : tensor<1xf64>
    %866 = stablehlo.subtract %860, %865 : tensor<1xf64>
    %867 = stablehlo.slice %813 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %868 = stablehlo.reshape %867 : (tensor<1x1xf64>) -> tensor<1xf64>
    %869 = stablehlo.slice %819 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %870 = stablehlo.reshape %869 : (tensor<1x1xf64>) -> tensor<1xf64>
    %871 = stablehlo.multiply %868, %870 : tensor<1xf64>
    %872 = stablehlo.slice %813 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %873 = stablehlo.reshape %872 : (tensor<1x1xf64>) -> tensor<1xf64>
    %874 = stablehlo.slice %819 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %875 = stablehlo.reshape %874 : (tensor<1x1xf64>) -> tensor<1xf64>
    %876 = stablehlo.multiply %873, %875 : tensor<1xf64>
    %877 = stablehlo.add %871, %876 : tensor<1xf64>
    %878 = stablehlo.slice %813 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %879 = stablehlo.reshape %878 : (tensor<1x1xf64>) -> tensor<1xf64>
    %880 = stablehlo.slice %819 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %881 = stablehlo.reshape %880 : (tensor<1x1xf64>) -> tensor<1xf64>
    %882 = stablehlo.multiply %879, %881 : tensor<1xf64>
    %883 = stablehlo.add %877, %882 : tensor<1xf64>
    %884 = stablehlo.slice %813 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %885 = stablehlo.reshape %884 : (tensor<1x1xf64>) -> tensor<1xf64>
    %886 = stablehlo.slice %819 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %887 = stablehlo.reshape %886 : (tensor<1x1xf64>) -> tensor<1xf64>
    %888 = stablehlo.multiply %885, %887 : tensor<1xf64>
    %889 = stablehlo.subtract %883, %888 : tensor<1xf64>
    %890 = stablehlo.slice %813 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %891 = stablehlo.reshape %890 : (tensor<1x1xf64>) -> tensor<1xf64>
    %892 = stablehlo.slice %819 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %893 = stablehlo.reshape %892 : (tensor<1x1xf64>) -> tensor<1xf64>
    %894 = stablehlo.multiply %891, %893 : tensor<1xf64>
    %895 = stablehlo.slice %813 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %896 = stablehlo.reshape %895 : (tensor<1x1xf64>) -> tensor<1xf64>
    %897 = stablehlo.slice %819 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %898 = stablehlo.reshape %897 : (tensor<1x1xf64>) -> tensor<1xf64>
    %899 = stablehlo.multiply %896, %898 : tensor<1xf64>
    %900 = stablehlo.subtract %894, %899 : tensor<1xf64>
    %901 = stablehlo.slice %813 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %902 = stablehlo.reshape %901 : (tensor<1x1xf64>) -> tensor<1xf64>
    %903 = stablehlo.slice %819 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %904 = stablehlo.reshape %903 : (tensor<1x1xf64>) -> tensor<1xf64>
    %905 = stablehlo.multiply %902, %904 : tensor<1xf64>
    %906 = stablehlo.add %900, %905 : tensor<1xf64>
    %907 = stablehlo.slice %813 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %908 = stablehlo.reshape %907 : (tensor<1x1xf64>) -> tensor<1xf64>
    %909 = stablehlo.slice %819 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %910 = stablehlo.reshape %909 : (tensor<1x1xf64>) -> tensor<1xf64>
    %911 = stablehlo.multiply %908, %910 : tensor<1xf64>
    %912 = stablehlo.add %906, %911 : tensor<1xf64>
    %913 = stablehlo.slice %813 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %914 = stablehlo.reshape %913 : (tensor<1x1xf64>) -> tensor<1xf64>
    %915 = stablehlo.slice %819 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %916 = stablehlo.reshape %915 : (tensor<1x1xf64>) -> tensor<1xf64>
    %917 = stablehlo.multiply %914, %916 : tensor<1xf64>
    %918 = stablehlo.slice %813 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %919 = stablehlo.reshape %918 : (tensor<1x1xf64>) -> tensor<1xf64>
    %920 = stablehlo.slice %819 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %921 = stablehlo.reshape %920 : (tensor<1x1xf64>) -> tensor<1xf64>
    %922 = stablehlo.multiply %919, %921 : tensor<1xf64>
    %923 = stablehlo.add %917, %922 : tensor<1xf64>
    %924 = stablehlo.slice %813 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %925 = stablehlo.reshape %924 : (tensor<1x1xf64>) -> tensor<1xf64>
    %926 = stablehlo.slice %819 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %927 = stablehlo.reshape %926 : (tensor<1x1xf64>) -> tensor<1xf64>
    %928 = stablehlo.multiply %925, %927 : tensor<1xf64>
    %929 = stablehlo.subtract %923, %928 : tensor<1xf64>
    %930 = stablehlo.slice %813 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %931 = stablehlo.reshape %930 : (tensor<1x1xf64>) -> tensor<1xf64>
    %932 = stablehlo.slice %819 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %933 = stablehlo.reshape %932 : (tensor<1x1xf64>) -> tensor<1xf64>
    %934 = stablehlo.multiply %931, %933 : tensor<1xf64>
    %935 = stablehlo.add %929, %934 : tensor<1xf64>
    %936 = stablehlo.broadcast_in_dim %866, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %937 = stablehlo.broadcast_in_dim %889, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %938 = stablehlo.broadcast_in_dim %912, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %939 = stablehlo.broadcast_in_dim %935, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %940 = stablehlo.concatenate %936, %937, %938, %939, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %941 = stablehlo.slice %814 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %942 = stablehlo.reshape %941 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %943 = stablehlo.slice %940 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %944 = stablehlo.reshape %943 : (tensor<1x1xf64>) -> tensor<1xf64>
    %945 = stablehlo.slice %940 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %946 = stablehlo.dot_general %945, %942, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %947 = stablehlo.broadcast_in_dim %946, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %948 = stablehlo.broadcast_in_dim %947, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %949 = stablehlo.multiply %948, %945 : tensor<1x3xf64>
    %cst_102 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %950 = stablehlo.broadcast_in_dim %cst_102, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %951 = stablehlo.multiply %950, %949 : tensor<1x3xf64>
    %952 = stablehlo.multiply %944, %944 : tensor<1xf64>
    %953 = stablehlo.dot_general %945, %945, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %954 = stablehlo.subtract %952, %953 : tensor<1xf64>
    %955 = stablehlo.broadcast_in_dim %954, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %956 = stablehlo.broadcast_in_dim %955, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %957 = stablehlo.multiply %956, %942 : tensor<1x3xf64>
    %958 = stablehlo.add %951, %957 : tensor<1x3xf64>
    %cst_103 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %959 = stablehlo.broadcast_in_dim %cst_103, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %960 = stablehlo.multiply %959, %944 : tensor<1xf64>
    %961 = call @cross(%945, %942) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %962 = stablehlo.broadcast_in_dim %960, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %963 = stablehlo.broadcast_in_dim %962, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %964 = stablehlo.multiply %963, %961 : tensor<1x3xf64>
    %965 = stablehlo.add %958, %964 : tensor<1x3xf64>
    %966 = stablehlo.add %965, %843 : tensor<1x3xf64>
    %967 = stablehlo.slice %815 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %968 = stablehlo.reshape %967 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %969 = stablehlo.slice %940 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %970 = stablehlo.reshape %969 : (tensor<1x1xf64>) -> tensor<1xf64>
    %971 = stablehlo.slice %940 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %972 = stablehlo.dot_general %971, %968, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %973 = stablehlo.broadcast_in_dim %972, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %974 = stablehlo.broadcast_in_dim %973, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %975 = stablehlo.multiply %974, %971 : tensor<1x3xf64>
    %cst_104 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %976 = stablehlo.broadcast_in_dim %cst_104, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %977 = stablehlo.multiply %976, %975 : tensor<1x3xf64>
    %978 = stablehlo.multiply %970, %970 : tensor<1xf64>
    %979 = stablehlo.dot_general %971, %971, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %980 = stablehlo.subtract %978, %979 : tensor<1xf64>
    %981 = stablehlo.broadcast_in_dim %980, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %982 = stablehlo.broadcast_in_dim %981, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %983 = stablehlo.multiply %982, %968 : tensor<1x3xf64>
    %984 = stablehlo.add %977, %983 : tensor<1x3xf64>
    %cst_105 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %985 = stablehlo.broadcast_in_dim %cst_105, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %986 = stablehlo.multiply %985, %970 : tensor<1xf64>
    %987 = call @cross(%971, %968) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %988 = stablehlo.broadcast_in_dim %986, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %989 = stablehlo.broadcast_in_dim %988, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %990 = stablehlo.multiply %989, %987 : tensor<1x3xf64>
    %991 = stablehlo.add %984, %990 : tensor<1x3xf64>
    %992 = stablehlo.slice %816 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %993 = stablehlo.reshape %992 : (tensor<1x1xf64>) -> tensor<1xf64>
    %994 = stablehlo.slice %817 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %995 = stablehlo.reshape %994 : (tensor<1x1xf64>) -> tensor<1xf64>
    %996 = stablehlo.subtract %993, %995 : tensor<1xf64>
    %997 = stablehlo.slice %815 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %998 = stablehlo.reshape %997 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %cst_106 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %999 = stablehlo.broadcast_in_dim %cst_106, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1000 = stablehlo.multiply %996, %999 : tensor<1xf64>
    %1001 = stablehlo.sine %1000 : tensor<1xf64>
    %cst_107 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %1002 = stablehlo.broadcast_in_dim %cst_107, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1003 = stablehlo.multiply %996, %1002 : tensor<1xf64>
    %1004 = stablehlo.cosine %1003 : tensor<1xf64>
    %1005 = stablehlo.broadcast_in_dim %1001, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1006 = stablehlo.broadcast_in_dim %1005, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1007 = stablehlo.multiply %998, %1006 : tensor<1x3xf64>
    %1008 = stablehlo.broadcast_in_dim %1004, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %c_108 = stablehlo.constant dense<0> : tensor<i64>
    %1009 = stablehlo.reshape %c_108 : (tensor<i64>) -> tensor<1xi64>
    %1010 = stablehlo.slice %1009 [0:1] : (tensor<1xi64>) -> tensor<1xi64>
    %1011 = stablehlo.reshape %1010 : (tensor<1xi64>) -> tensor<i64>
    %1012 = stablehlo.broadcast_in_dim %1011, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %cst_109 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %1013 = stablehlo.broadcast_in_dim %cst_109, dims = [] : (tensor<f64>) -> tensor<4xf64>
    %c_110 = stablehlo.constant dense<0> : tensor<i64>
    %1014 = stablehlo.broadcast_in_dim %c_110, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1015 = stablehlo.compare  LT, %1012, %1014,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_111 = stablehlo.constant dense<3> : tensor<i64>
    %1016 = stablehlo.broadcast_in_dim %c_111, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1017 = stablehlo.add %1012, %1016 : tensor<1xi64>
    %1018 = call @_where_7(%1015, %1017, %1012) : (tensor<1xi1>, tensor<1xi64>, tensor<1xi64>) -> tensor<1xi64>
    %c_112 = stablehlo.constant dense<0> : tensor<i64>
    %c_113 = stablehlo.constant dense<3> : tensor<i64>
    %1019 = call @clip(%1018, %c_112, %c_113) : (tensor<1xi64>, tensor<i64>, tensor<i64>) -> tensor<1xi64>
    %1020 = call @argsort(%1019) : (tensor<1xi64>) -> tensor<1xi64>
    %1021 = stablehlo.iota dim = 0 : tensor<1xi64>
    %c_114 = stablehlo.constant dense<0> : tensor<i64>
    %1022 = stablehlo.broadcast_in_dim %c_114, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1023 = stablehlo.compare  LT, %1020, %1022,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_115 = stablehlo.constant dense<1> : tensor<i64>
    %1024 = stablehlo.broadcast_in_dim %c_115, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1025 = stablehlo.add %1020, %1024 : tensor<1xi64>
    %1026 = stablehlo.select %1023, %1025, %1020 : tensor<1xi1>, tensor<1xi64>
    %1027 = stablehlo.convert %1026 : (tensor<1xi64>) -> tensor<1xi32>
    %1028 = stablehlo.broadcast_in_dim %1027, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %1029 = stablehlo.convert %1019 : tensor<1xi64>
    %1030 = "stablehlo.scatter"(%1029, %1028, %1021) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<1xi64>, tensor<1x1xi32>, tensor<1xi64>) -> tensor<1xi64>
    %1031 = stablehlo.convert %1030 : tensor<1xi64>
    %c_116 = stablehlo.constant dense<true> : tensor<i1>
    %1032 = stablehlo.broadcast_in_dim %c_116, dims = [] : (tensor<i1>) -> tensor<4xi1>
    %c_117 = stablehlo.constant dense<0> : tensor<i64>
    %1033 = stablehlo.broadcast_in_dim %c_117, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1034 = stablehlo.compare  LT, %1031, %1033,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_118 = stablehlo.constant dense<4> : tensor<i64>
    %1035 = stablehlo.broadcast_in_dim %c_118, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1036 = stablehlo.add %1031, %1035 : tensor<1xi64>
    %1037 = stablehlo.select %1034, %1036, %1031 : tensor<1xi1>, tensor<1xi64>
    %1038 = stablehlo.convert %1037 : (tensor<1xi64>) -> tensor<1xi32>
    %1039 = stablehlo.broadcast_in_dim %1038, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %c_119 = stablehlo.constant dense<false> : tensor<i1>
    %1040 = stablehlo.broadcast_in_dim %c_119, dims = [] : (tensor<i1>) -> tensor<1xi1>
    %1041 = "stablehlo.scatter"(%1032, %1039, %1040) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i1>, %arg114: tensor<i1>):
      stablehlo.return %arg114 : tensor<i1>
    }) : (tensor<4xi1>, tensor<1x1xi32>, tensor<1xi1>) -> tensor<4xi1>
    %1042 = call @_cumulative_reduction(%1041) : (tensor<4xi1>) -> tensor<4xi64>
    %c_120 = stablehlo.constant dense<0> : tensor<i64>
    %1043 = stablehlo.broadcast_in_dim %c_120, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %c_121 = stablehlo.constant dense<0> : tensor<i64>
    %1044 = call @clip_8(%1042, %c_121) : (tensor<4xi64>, tensor<i64>) -> tensor<4xi64>
    %c_122 = stablehlo.constant dense<1> : tensor<i64>
    %c_123 = stablehlo.constant dense<0> : tensor<i64>
    %1045 = stablehlo.broadcast_in_dim %c_123, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %1046 = stablehlo.compare  LT, %1044, %1045,  SIGNED : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1>
    %c_124 = stablehlo.constant dense<3> : tensor<i64>
    %1047 = stablehlo.broadcast_in_dim %c_124, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %1048 = stablehlo.add %1044, %1047 : tensor<4xi64>
    %1049 = stablehlo.select %1046, %1048, %1044 : tensor<4xi1>, tensor<4xi64>
    %1050 = stablehlo.convert %1049 : (tensor<4xi64>) -> tensor<4xi32>
    %1051 = stablehlo.broadcast_in_dim %1050, dims = [0] : (tensor<4xi32>) -> tensor<4x1xi32>
    %1052 = stablehlo.broadcast_in_dim %c_122, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %1053 = "stablehlo.scatter"(%1043, %1051, %1052) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<3xi64>, tensor<4x1xi32>, tensor<4xi64>) -> tensor<3xi64>
    %1054 = call @_cumulative_reduction_9(%1053) : (tensor<3xi64>) -> tensor<3xi64>
    %c_125 = stablehlo.constant dense<1> : tensor<i64>
    %1055 = call @floor_divide(%1054, %c_125) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_126 = stablehlo.constant dense<4> : tensor<i64>
    %1056 = call @remainder_12(%1055, %c_126) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_127 = stablehlo.constant dense<0> : tensor<i64>
    %1057 = stablehlo.broadcast_in_dim %c_127, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1058 = stablehlo.compare  LT, %1031, %1057,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_128 = stablehlo.constant dense<4> : tensor<i64>
    %1059 = stablehlo.broadcast_in_dim %c_128, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1060 = stablehlo.add %1031, %1059 : tensor<1xi64>
    %1061 = stablehlo.select %1058, %1060, %1031 : tensor<1xi1>, tensor<1xi64>
    %1062 = stablehlo.convert %1061 : (tensor<1xi64>) -> tensor<1xi32>
    %1063 = stablehlo.broadcast_in_dim %1062, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %1064 = stablehlo.broadcast_in_dim %1013, dims = [1] : (tensor<4xf64>) -> tensor<1x4xf64>
    %1065 = "stablehlo.scatter"(%1064, %1063, %1008) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<1x1xi32>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %c_129 = stablehlo.constant dense<0> : tensor<i64>
    %1066 = stablehlo.broadcast_in_dim %c_129, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %1067 = stablehlo.compare  LT, %1056, %1066,  SIGNED : (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi1>
    %c_130 = stablehlo.constant dense<4> : tensor<i64>
    %1068 = stablehlo.broadcast_in_dim %c_130, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %1069 = stablehlo.add %1056, %1068 : tensor<3xi64>
    %1070 = stablehlo.select %1067, %1069, %1056 : tensor<3xi1>, tensor<3xi64>
    %1071 = stablehlo.convert %1070 : (tensor<3xi64>) -> tensor<3xi32>
    %1072 = stablehlo.broadcast_in_dim %1071, dims = [0] : (tensor<3xi32>) -> tensor<3x1xi32>
    %1073 = "stablehlo.scatter"(%1065, %1072, %1007) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<3x1xi32>, tensor<1x3xf64>) -> tensor<1x4xf64>
    %1074 = stablehlo.slice %940 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1075 = stablehlo.reshape %1074 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1076 = stablehlo.slice %1073 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1077 = stablehlo.reshape %1076 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1078 = stablehlo.multiply %1075, %1077 : tensor<1xf64>
    %1079 = stablehlo.slice %940 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1080 = stablehlo.reshape %1079 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1081 = stablehlo.slice %1073 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1082 = stablehlo.reshape %1081 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1083 = stablehlo.multiply %1080, %1082 : tensor<1xf64>
    %1084 = stablehlo.subtract %1078, %1083 : tensor<1xf64>
    %1085 = stablehlo.slice %940 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1086 = stablehlo.reshape %1085 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1087 = stablehlo.slice %1073 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1088 = stablehlo.reshape %1087 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1089 = stablehlo.multiply %1086, %1088 : tensor<1xf64>
    %1090 = stablehlo.subtract %1084, %1089 : tensor<1xf64>
    %1091 = stablehlo.slice %940 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1092 = stablehlo.reshape %1091 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1093 = stablehlo.slice %1073 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1094 = stablehlo.reshape %1093 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1095 = stablehlo.multiply %1092, %1094 : tensor<1xf64>
    %1096 = stablehlo.subtract %1090, %1095 : tensor<1xf64>
    %1097 = stablehlo.slice %940 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1098 = stablehlo.reshape %1097 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1099 = stablehlo.slice %1073 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1100 = stablehlo.reshape %1099 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1101 = stablehlo.multiply %1098, %1100 : tensor<1xf64>
    %1102 = stablehlo.slice %940 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1103 = stablehlo.reshape %1102 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1104 = stablehlo.slice %1073 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1105 = stablehlo.reshape %1104 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1106 = stablehlo.multiply %1103, %1105 : tensor<1xf64>
    %1107 = stablehlo.add %1101, %1106 : tensor<1xf64>
    %1108 = stablehlo.slice %940 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1109 = stablehlo.reshape %1108 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1110 = stablehlo.slice %1073 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1111 = stablehlo.reshape %1110 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1112 = stablehlo.multiply %1109, %1111 : tensor<1xf64>
    %1113 = stablehlo.add %1107, %1112 : tensor<1xf64>
    %1114 = stablehlo.slice %940 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1115 = stablehlo.reshape %1114 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1116 = stablehlo.slice %1073 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1117 = stablehlo.reshape %1116 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1118 = stablehlo.multiply %1115, %1117 : tensor<1xf64>
    %1119 = stablehlo.subtract %1113, %1118 : tensor<1xf64>
    %1120 = stablehlo.slice %940 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1121 = stablehlo.reshape %1120 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1122 = stablehlo.slice %1073 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1123 = stablehlo.reshape %1122 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1124 = stablehlo.multiply %1121, %1123 : tensor<1xf64>
    %1125 = stablehlo.slice %940 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1126 = stablehlo.reshape %1125 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1127 = stablehlo.slice %1073 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1128 = stablehlo.reshape %1127 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1129 = stablehlo.multiply %1126, %1128 : tensor<1xf64>
    %1130 = stablehlo.subtract %1124, %1129 : tensor<1xf64>
    %1131 = stablehlo.slice %940 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1132 = stablehlo.reshape %1131 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1133 = stablehlo.slice %1073 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1134 = stablehlo.reshape %1133 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1135 = stablehlo.multiply %1132, %1134 : tensor<1xf64>
    %1136 = stablehlo.add %1130, %1135 : tensor<1xf64>
    %1137 = stablehlo.slice %940 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1138 = stablehlo.reshape %1137 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1139 = stablehlo.slice %1073 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1140 = stablehlo.reshape %1139 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1141 = stablehlo.multiply %1138, %1140 : tensor<1xf64>
    %1142 = stablehlo.add %1136, %1141 : tensor<1xf64>
    %1143 = stablehlo.slice %940 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1144 = stablehlo.reshape %1143 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1145 = stablehlo.slice %1073 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1146 = stablehlo.reshape %1145 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1147 = stablehlo.multiply %1144, %1146 : tensor<1xf64>
    %1148 = stablehlo.slice %940 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1149 = stablehlo.reshape %1148 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1150 = stablehlo.slice %1073 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1151 = stablehlo.reshape %1150 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1152 = stablehlo.multiply %1149, %1151 : tensor<1xf64>
    %1153 = stablehlo.add %1147, %1152 : tensor<1xf64>
    %1154 = stablehlo.slice %940 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1155 = stablehlo.reshape %1154 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1156 = stablehlo.slice %1073 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1157 = stablehlo.reshape %1156 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1158 = stablehlo.multiply %1155, %1157 : tensor<1xf64>
    %1159 = stablehlo.subtract %1153, %1158 : tensor<1xf64>
    %1160 = stablehlo.slice %940 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1161 = stablehlo.reshape %1160 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1162 = stablehlo.slice %1073 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1163 = stablehlo.reshape %1162 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1164 = stablehlo.multiply %1161, %1163 : tensor<1xf64>
    %1165 = stablehlo.add %1159, %1164 : tensor<1xf64>
    %1166 = stablehlo.broadcast_in_dim %1096, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1167 = stablehlo.broadcast_in_dim %1119, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1168 = stablehlo.broadcast_in_dim %1142, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1169 = stablehlo.broadcast_in_dim %1165, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1170 = stablehlo.concatenate %1166, %1167, %1168, %1169, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %1171 = stablehlo.slice %814 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %1172 = stablehlo.reshape %1171 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %1173 = stablehlo.slice %1170 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1174 = stablehlo.reshape %1173 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1175 = stablehlo.slice %1170 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %1176 = stablehlo.dot_general %1175, %1172, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1177 = stablehlo.broadcast_in_dim %1176, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1178 = stablehlo.broadcast_in_dim %1177, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1179 = stablehlo.multiply %1178, %1175 : tensor<1x3xf64>
    %cst_131 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1180 = stablehlo.broadcast_in_dim %cst_131, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %1181 = stablehlo.multiply %1180, %1179 : tensor<1x3xf64>
    %1182 = stablehlo.multiply %1174, %1174 : tensor<1xf64>
    %1183 = stablehlo.dot_general %1175, %1175, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1184 = stablehlo.subtract %1182, %1183 : tensor<1xf64>
    %1185 = stablehlo.broadcast_in_dim %1184, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1186 = stablehlo.broadcast_in_dim %1185, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1187 = stablehlo.multiply %1186, %1172 : tensor<1x3xf64>
    %1188 = stablehlo.add %1181, %1187 : tensor<1x3xf64>
    %cst_132 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1189 = stablehlo.broadcast_in_dim %cst_132, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1190 = stablehlo.multiply %1189, %1174 : tensor<1xf64>
    %1191 = call @cross(%1175, %1172) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %1192 = stablehlo.broadcast_in_dim %1190, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1193 = stablehlo.broadcast_in_dim %1192, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1194 = stablehlo.multiply %1193, %1191 : tensor<1x3xf64>
    %1195 = stablehlo.add %1188, %1194 : tensor<1x3xf64>
    %1196 = stablehlo.subtract %966, %1195 : tensor<1x3xf64>
    %1197 = stablehlo.broadcast_in_dim %966, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %1198 = stablehlo.broadcast_in_dim %991, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %1199 = stablehlo.broadcast_in_dim %1170, dims = [0, 1] : (tensor<1x4xf64>) -> tensor<1x4x1xf64>
    %1200 = stablehlo.broadcast_in_dim %1170, dims = [0, 2] : (tensor<1x4xf64>) -> tensor<1x1x4xf64>
    %1201 = stablehlo.broadcast_in_dim %1199, dims = [0, 1, 2] : (tensor<1x4x1xf64>) -> tensor<1x4x4xf64>
    %1202 = stablehlo.broadcast_in_dim %1200, dims = [0, 1, 2] : (tensor<1x1x4xf64>) -> tensor<1x4x4xf64>
    %1203 = stablehlo.multiply %1201, %1202 : tensor<1x4x4xf64>
    %1204 = stablehlo.slice %1203 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1205 = stablehlo.reshape %1204 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1206 = stablehlo.slice %1203 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1207 = stablehlo.reshape %1206 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1208 = stablehlo.add %1205, %1207 : tensor<1xf64>
    %1209 = stablehlo.slice %1203 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1210 = stablehlo.reshape %1209 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1211 = stablehlo.subtract %1208, %1210 : tensor<1xf64>
    %1212 = stablehlo.slice %1203 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1213 = stablehlo.reshape %1212 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1214 = stablehlo.subtract %1211, %1213 : tensor<1xf64>
    %1215 = stablehlo.slice %1203 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1216 = stablehlo.reshape %1215 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1217 = stablehlo.slice %1203 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1218 = stablehlo.reshape %1217 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1219 = stablehlo.subtract %1216, %1218 : tensor<1xf64>
    %cst_133 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1220 = stablehlo.broadcast_in_dim %cst_133, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1221 = stablehlo.multiply %1220, %1219 : tensor<1xf64>
    %1222 = stablehlo.slice %1203 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1223 = stablehlo.reshape %1222 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1224 = stablehlo.slice %1203 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1225 = stablehlo.reshape %1224 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1226 = stablehlo.add %1223, %1225 : tensor<1xf64>
    %cst_134 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1227 = stablehlo.broadcast_in_dim %cst_134, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1228 = stablehlo.multiply %1227, %1226 : tensor<1xf64>
    %1229 = stablehlo.slice %1203 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1230 = stablehlo.reshape %1229 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1231 = stablehlo.slice %1203 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1232 = stablehlo.reshape %1231 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1233 = stablehlo.add %1230, %1232 : tensor<1xf64>
    %cst_135 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1234 = stablehlo.broadcast_in_dim %cst_135, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1235 = stablehlo.multiply %1234, %1233 : tensor<1xf64>
    %1236 = stablehlo.slice %1203 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1237 = stablehlo.reshape %1236 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1238 = stablehlo.slice %1203 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1239 = stablehlo.reshape %1238 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1240 = stablehlo.subtract %1237, %1239 : tensor<1xf64>
    %1241 = stablehlo.slice %1203 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1242 = stablehlo.reshape %1241 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1243 = stablehlo.add %1240, %1242 : tensor<1xf64>
    %1244 = stablehlo.slice %1203 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1245 = stablehlo.reshape %1244 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1246 = stablehlo.subtract %1243, %1245 : tensor<1xf64>
    %1247 = stablehlo.slice %1203 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1248 = stablehlo.reshape %1247 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1249 = stablehlo.slice %1203 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1250 = stablehlo.reshape %1249 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1251 = stablehlo.subtract %1248, %1250 : tensor<1xf64>
    %cst_136 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1252 = stablehlo.broadcast_in_dim %cst_136, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1253 = stablehlo.multiply %1252, %1251 : tensor<1xf64>
    %1254 = stablehlo.slice %1203 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1255 = stablehlo.reshape %1254 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1256 = stablehlo.slice %1203 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1257 = stablehlo.reshape %1256 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1258 = stablehlo.subtract %1255, %1257 : tensor<1xf64>
    %cst_137 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1259 = stablehlo.broadcast_in_dim %cst_137, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1260 = stablehlo.multiply %1259, %1258 : tensor<1xf64>
    %1261 = stablehlo.slice %1203 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1262 = stablehlo.reshape %1261 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1263 = stablehlo.slice %1203 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1264 = stablehlo.reshape %1263 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1265 = stablehlo.add %1262, %1264 : tensor<1xf64>
    %cst_138 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1266 = stablehlo.broadcast_in_dim %cst_138, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1267 = stablehlo.multiply %1266, %1265 : tensor<1xf64>
    %1268 = stablehlo.slice %1203 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1269 = stablehlo.reshape %1268 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1270 = stablehlo.slice %1203 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1271 = stablehlo.reshape %1270 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1272 = stablehlo.subtract %1269, %1271 : tensor<1xf64>
    %1273 = stablehlo.slice %1203 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1274 = stablehlo.reshape %1273 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1275 = stablehlo.subtract %1272, %1274 : tensor<1xf64>
    %1276 = stablehlo.slice %1203 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1277 = stablehlo.reshape %1276 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1278 = stablehlo.add %1275, %1277 : tensor<1xf64>
    %1279 = stablehlo.broadcast_in_dim %1214, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1280 = stablehlo.broadcast_in_dim %1221, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1281 = stablehlo.broadcast_in_dim %1228, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1282 = stablehlo.concatenate %1279, %1280, %1281, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %1283 = stablehlo.broadcast_in_dim %1235, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1284 = stablehlo.broadcast_in_dim %1246, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1285 = stablehlo.broadcast_in_dim %1253, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1286 = stablehlo.concatenate %1283, %1284, %1285, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %1287 = stablehlo.broadcast_in_dim %1260, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1288 = stablehlo.broadcast_in_dim %1267, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1289 = stablehlo.broadcast_in_dim %1278, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1290 = stablehlo.concatenate %1287, %1288, %1289, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %1291 = stablehlo.broadcast_in_dim %1282, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %1292 = stablehlo.broadcast_in_dim %1286, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %1293 = stablehlo.broadcast_in_dim %1290, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %1294 = stablehlo.concatenate %1291, %1292, %1293, dim = 1 : (tensor<1x1x3xf64>, tensor<1x1x3xf64>, tensor<1x1x3xf64>) -> tensor<1x3x3xf64>
    %1295 = call @_take_2(%1196, %c_11) : (tensor<1x3xf64>, tensor<1xi64>) -> tensor<1x3xf64>
    %1296 = call @_take_3(%1170, %c_11) : (tensor<1x4xf64>, tensor<1xi64>) -> tensor<1x4xf64>
    %1297 = call @_take_4(%arg5, %c_12) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %1298 = call @_take_4(%arg6, %c_13) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %1299 = call @_take_6(%arg13, %c_14) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %1300 = call @_take_6(%arg0, %c_15) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %1301 = stablehlo.slice %arg1 [4:5, 0:3] : (tensor<9x3xf64>) -> tensor<1x3xf64>
    %1302 = stablehlo.slice %arg2 [4:5, 0:4] : (tensor<9x4xf64>) -> tensor<1x4xf64>
    %1303 = stablehlo.slice %1296 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1304 = stablehlo.reshape %1303 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1305 = stablehlo.slice %1296 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %1306 = stablehlo.dot_general %1305, %1301, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1307 = stablehlo.broadcast_in_dim %1306, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1308 = stablehlo.broadcast_in_dim %1307, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1309 = stablehlo.multiply %1308, %1305 : tensor<1x3xf64>
    %cst_139 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1310 = stablehlo.broadcast_in_dim %cst_139, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %1311 = stablehlo.multiply %1310, %1309 : tensor<1x3xf64>
    %1312 = stablehlo.multiply %1304, %1304 : tensor<1xf64>
    %1313 = stablehlo.dot_general %1305, %1305, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1314 = stablehlo.subtract %1312, %1313 : tensor<1xf64>
    %1315 = stablehlo.broadcast_in_dim %1314, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1316 = stablehlo.broadcast_in_dim %1315, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1317 = stablehlo.multiply %1316, %1301 : tensor<1x3xf64>
    %1318 = stablehlo.add %1311, %1317 : tensor<1x3xf64>
    %cst_140 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1319 = stablehlo.broadcast_in_dim %cst_140, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1320 = stablehlo.multiply %1319, %1304 : tensor<1xf64>
    %1321 = call @cross(%1305, %1301) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %1322 = stablehlo.broadcast_in_dim %1320, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1323 = stablehlo.broadcast_in_dim %1322, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1324 = stablehlo.multiply %1323, %1321 : tensor<1x3xf64>
    %1325 = stablehlo.add %1318, %1324 : tensor<1x3xf64>
    %1326 = stablehlo.add %1295, %1325 : tensor<1x3xf64>
    %1327 = stablehlo.slice %1296 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1328 = stablehlo.reshape %1327 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1329 = stablehlo.slice %1302 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1330 = stablehlo.reshape %1329 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1331 = stablehlo.multiply %1328, %1330 : tensor<1xf64>
    %1332 = stablehlo.slice %1296 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1333 = stablehlo.reshape %1332 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1334 = stablehlo.slice %1302 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1335 = stablehlo.reshape %1334 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1336 = stablehlo.multiply %1333, %1335 : tensor<1xf64>
    %1337 = stablehlo.subtract %1331, %1336 : tensor<1xf64>
    %1338 = stablehlo.slice %1296 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1339 = stablehlo.reshape %1338 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1340 = stablehlo.slice %1302 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1341 = stablehlo.reshape %1340 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1342 = stablehlo.multiply %1339, %1341 : tensor<1xf64>
    %1343 = stablehlo.subtract %1337, %1342 : tensor<1xf64>
    %1344 = stablehlo.slice %1296 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1345 = stablehlo.reshape %1344 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1346 = stablehlo.slice %1302 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1347 = stablehlo.reshape %1346 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1348 = stablehlo.multiply %1345, %1347 : tensor<1xf64>
    %1349 = stablehlo.subtract %1343, %1348 : tensor<1xf64>
    %1350 = stablehlo.slice %1296 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1351 = stablehlo.reshape %1350 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1352 = stablehlo.slice %1302 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1353 = stablehlo.reshape %1352 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1354 = stablehlo.multiply %1351, %1353 : tensor<1xf64>
    %1355 = stablehlo.slice %1296 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1356 = stablehlo.reshape %1355 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1357 = stablehlo.slice %1302 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1358 = stablehlo.reshape %1357 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1359 = stablehlo.multiply %1356, %1358 : tensor<1xf64>
    %1360 = stablehlo.add %1354, %1359 : tensor<1xf64>
    %1361 = stablehlo.slice %1296 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1362 = stablehlo.reshape %1361 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1363 = stablehlo.slice %1302 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1364 = stablehlo.reshape %1363 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1365 = stablehlo.multiply %1362, %1364 : tensor<1xf64>
    %1366 = stablehlo.add %1360, %1365 : tensor<1xf64>
    %1367 = stablehlo.slice %1296 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1368 = stablehlo.reshape %1367 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1369 = stablehlo.slice %1302 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1370 = stablehlo.reshape %1369 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1371 = stablehlo.multiply %1368, %1370 : tensor<1xf64>
    %1372 = stablehlo.subtract %1366, %1371 : tensor<1xf64>
    %1373 = stablehlo.slice %1296 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1374 = stablehlo.reshape %1373 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1375 = stablehlo.slice %1302 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1376 = stablehlo.reshape %1375 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1377 = stablehlo.multiply %1374, %1376 : tensor<1xf64>
    %1378 = stablehlo.slice %1296 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1379 = stablehlo.reshape %1378 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1380 = stablehlo.slice %1302 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1381 = stablehlo.reshape %1380 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1382 = stablehlo.multiply %1379, %1381 : tensor<1xf64>
    %1383 = stablehlo.subtract %1377, %1382 : tensor<1xf64>
    %1384 = stablehlo.slice %1296 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1385 = stablehlo.reshape %1384 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1386 = stablehlo.slice %1302 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1387 = stablehlo.reshape %1386 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1388 = stablehlo.multiply %1385, %1387 : tensor<1xf64>
    %1389 = stablehlo.add %1383, %1388 : tensor<1xf64>
    %1390 = stablehlo.slice %1296 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1391 = stablehlo.reshape %1390 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1392 = stablehlo.slice %1302 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1393 = stablehlo.reshape %1392 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1394 = stablehlo.multiply %1391, %1393 : tensor<1xf64>
    %1395 = stablehlo.add %1389, %1394 : tensor<1xf64>
    %1396 = stablehlo.slice %1296 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1397 = stablehlo.reshape %1396 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1398 = stablehlo.slice %1302 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1399 = stablehlo.reshape %1398 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1400 = stablehlo.multiply %1397, %1399 : tensor<1xf64>
    %1401 = stablehlo.slice %1296 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1402 = stablehlo.reshape %1401 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1403 = stablehlo.slice %1302 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1404 = stablehlo.reshape %1403 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1405 = stablehlo.multiply %1402, %1404 : tensor<1xf64>
    %1406 = stablehlo.add %1400, %1405 : tensor<1xf64>
    %1407 = stablehlo.slice %1296 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1408 = stablehlo.reshape %1407 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1409 = stablehlo.slice %1302 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1410 = stablehlo.reshape %1409 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1411 = stablehlo.multiply %1408, %1410 : tensor<1xf64>
    %1412 = stablehlo.subtract %1406, %1411 : tensor<1xf64>
    %1413 = stablehlo.slice %1296 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1414 = stablehlo.reshape %1413 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1415 = stablehlo.slice %1302 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1416 = stablehlo.reshape %1415 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1417 = stablehlo.multiply %1414, %1416 : tensor<1xf64>
    %1418 = stablehlo.add %1412, %1417 : tensor<1xf64>
    %1419 = stablehlo.broadcast_in_dim %1349, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1420 = stablehlo.broadcast_in_dim %1372, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1421 = stablehlo.broadcast_in_dim %1395, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1422 = stablehlo.broadcast_in_dim %1418, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1423 = stablehlo.concatenate %1419, %1420, %1421, %1422, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %1424 = stablehlo.slice %1297 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %1425 = stablehlo.reshape %1424 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %1426 = stablehlo.slice %1423 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1427 = stablehlo.reshape %1426 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1428 = stablehlo.slice %1423 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %1429 = stablehlo.dot_general %1428, %1425, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1430 = stablehlo.broadcast_in_dim %1429, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1431 = stablehlo.broadcast_in_dim %1430, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1432 = stablehlo.multiply %1431, %1428 : tensor<1x3xf64>
    %cst_141 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1433 = stablehlo.broadcast_in_dim %cst_141, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %1434 = stablehlo.multiply %1433, %1432 : tensor<1x3xf64>
    %1435 = stablehlo.multiply %1427, %1427 : tensor<1xf64>
    %1436 = stablehlo.dot_general %1428, %1428, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1437 = stablehlo.subtract %1435, %1436 : tensor<1xf64>
    %1438 = stablehlo.broadcast_in_dim %1437, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1439 = stablehlo.broadcast_in_dim %1438, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1440 = stablehlo.multiply %1439, %1425 : tensor<1x3xf64>
    %1441 = stablehlo.add %1434, %1440 : tensor<1x3xf64>
    %cst_142 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1442 = stablehlo.broadcast_in_dim %cst_142, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1443 = stablehlo.multiply %1442, %1427 : tensor<1xf64>
    %1444 = call @cross(%1428, %1425) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %1445 = stablehlo.broadcast_in_dim %1443, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1446 = stablehlo.broadcast_in_dim %1445, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1447 = stablehlo.multiply %1446, %1444 : tensor<1x3xf64>
    %1448 = stablehlo.add %1441, %1447 : tensor<1x3xf64>
    %1449 = stablehlo.add %1448, %1326 : tensor<1x3xf64>
    %1450 = stablehlo.slice %1298 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %1451 = stablehlo.reshape %1450 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %1452 = stablehlo.slice %1423 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1453 = stablehlo.reshape %1452 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1454 = stablehlo.slice %1423 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %1455 = stablehlo.dot_general %1454, %1451, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1456 = stablehlo.broadcast_in_dim %1455, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1457 = stablehlo.broadcast_in_dim %1456, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1458 = stablehlo.multiply %1457, %1454 : tensor<1x3xf64>
    %cst_143 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1459 = stablehlo.broadcast_in_dim %cst_143, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %1460 = stablehlo.multiply %1459, %1458 : tensor<1x3xf64>
    %1461 = stablehlo.multiply %1453, %1453 : tensor<1xf64>
    %1462 = stablehlo.dot_general %1454, %1454, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1463 = stablehlo.subtract %1461, %1462 : tensor<1xf64>
    %1464 = stablehlo.broadcast_in_dim %1463, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1465 = stablehlo.broadcast_in_dim %1464, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1466 = stablehlo.multiply %1465, %1451 : tensor<1x3xf64>
    %1467 = stablehlo.add %1460, %1466 : tensor<1x3xf64>
    %cst_144 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1468 = stablehlo.broadcast_in_dim %cst_144, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1469 = stablehlo.multiply %1468, %1453 : tensor<1xf64>
    %1470 = call @cross(%1454, %1451) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %1471 = stablehlo.broadcast_in_dim %1469, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1472 = stablehlo.broadcast_in_dim %1471, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1473 = stablehlo.multiply %1472, %1470 : tensor<1x3xf64>
    %1474 = stablehlo.add %1467, %1473 : tensor<1x3xf64>
    %1475 = stablehlo.slice %1299 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %1476 = stablehlo.reshape %1475 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1477 = stablehlo.slice %1300 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %1478 = stablehlo.reshape %1477 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1479 = stablehlo.subtract %1476, %1478 : tensor<1xf64>
    %1480 = stablehlo.slice %1298 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %1481 = stablehlo.reshape %1480 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %cst_145 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %1482 = stablehlo.broadcast_in_dim %cst_145, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1483 = stablehlo.multiply %1479, %1482 : tensor<1xf64>
    %1484 = stablehlo.sine %1483 : tensor<1xf64>
    %cst_146 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %1485 = stablehlo.broadcast_in_dim %cst_146, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1486 = stablehlo.multiply %1479, %1485 : tensor<1xf64>
    %1487 = stablehlo.cosine %1486 : tensor<1xf64>
    %1488 = stablehlo.broadcast_in_dim %1484, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1489 = stablehlo.broadcast_in_dim %1488, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1490 = stablehlo.multiply %1481, %1489 : tensor<1x3xf64>
    %1491 = stablehlo.broadcast_in_dim %1487, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %c_147 = stablehlo.constant dense<0> : tensor<i64>
    %1492 = stablehlo.reshape %c_147 : (tensor<i64>) -> tensor<1xi64>
    %1493 = stablehlo.slice %1492 [0:1] : (tensor<1xi64>) -> tensor<1xi64>
    %1494 = stablehlo.reshape %1493 : (tensor<1xi64>) -> tensor<i64>
    %1495 = stablehlo.broadcast_in_dim %1494, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %cst_148 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %1496 = stablehlo.broadcast_in_dim %cst_148, dims = [] : (tensor<f64>) -> tensor<4xf64>
    %c_149 = stablehlo.constant dense<0> : tensor<i64>
    %1497 = stablehlo.broadcast_in_dim %c_149, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1498 = stablehlo.compare  LT, %1495, %1497,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_150 = stablehlo.constant dense<3> : tensor<i64>
    %1499 = stablehlo.broadcast_in_dim %c_150, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1500 = stablehlo.add %1495, %1499 : tensor<1xi64>
    %1501 = call @_where_7(%1498, %1500, %1495) : (tensor<1xi1>, tensor<1xi64>, tensor<1xi64>) -> tensor<1xi64>
    %c_151 = stablehlo.constant dense<0> : tensor<i64>
    %c_152 = stablehlo.constant dense<3> : tensor<i64>
    %1502 = call @clip(%1501, %c_151, %c_152) : (tensor<1xi64>, tensor<i64>, tensor<i64>) -> tensor<1xi64>
    %1503 = call @argsort(%1502) : (tensor<1xi64>) -> tensor<1xi64>
    %1504 = stablehlo.iota dim = 0 : tensor<1xi64>
    %c_153 = stablehlo.constant dense<0> : tensor<i64>
    %1505 = stablehlo.broadcast_in_dim %c_153, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1506 = stablehlo.compare  LT, %1503, %1505,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_154 = stablehlo.constant dense<1> : tensor<i64>
    %1507 = stablehlo.broadcast_in_dim %c_154, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1508 = stablehlo.add %1503, %1507 : tensor<1xi64>
    %1509 = stablehlo.select %1506, %1508, %1503 : tensor<1xi1>, tensor<1xi64>
    %1510 = stablehlo.convert %1509 : (tensor<1xi64>) -> tensor<1xi32>
    %1511 = stablehlo.broadcast_in_dim %1510, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %1512 = stablehlo.convert %1502 : tensor<1xi64>
    %1513 = "stablehlo.scatter"(%1512, %1511, %1504) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<1xi64>, tensor<1x1xi32>, tensor<1xi64>) -> tensor<1xi64>
    %1514 = stablehlo.convert %1513 : tensor<1xi64>
    %c_155 = stablehlo.constant dense<true> : tensor<i1>
    %1515 = stablehlo.broadcast_in_dim %c_155, dims = [] : (tensor<i1>) -> tensor<4xi1>
    %c_156 = stablehlo.constant dense<0> : tensor<i64>
    %1516 = stablehlo.broadcast_in_dim %c_156, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1517 = stablehlo.compare  LT, %1514, %1516,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_157 = stablehlo.constant dense<4> : tensor<i64>
    %1518 = stablehlo.broadcast_in_dim %c_157, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1519 = stablehlo.add %1514, %1518 : tensor<1xi64>
    %1520 = stablehlo.select %1517, %1519, %1514 : tensor<1xi1>, tensor<1xi64>
    %1521 = stablehlo.convert %1520 : (tensor<1xi64>) -> tensor<1xi32>
    %1522 = stablehlo.broadcast_in_dim %1521, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %c_158 = stablehlo.constant dense<false> : tensor<i1>
    %1523 = stablehlo.broadcast_in_dim %c_158, dims = [] : (tensor<i1>) -> tensor<1xi1>
    %1524 = "stablehlo.scatter"(%1515, %1522, %1523) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i1>, %arg114: tensor<i1>):
      stablehlo.return %arg114 : tensor<i1>
    }) : (tensor<4xi1>, tensor<1x1xi32>, tensor<1xi1>) -> tensor<4xi1>
    %1525 = call @_cumulative_reduction(%1524) : (tensor<4xi1>) -> tensor<4xi64>
    %c_159 = stablehlo.constant dense<0> : tensor<i64>
    %1526 = stablehlo.broadcast_in_dim %c_159, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %c_160 = stablehlo.constant dense<0> : tensor<i64>
    %1527 = call @clip_8(%1525, %c_160) : (tensor<4xi64>, tensor<i64>) -> tensor<4xi64>
    %c_161 = stablehlo.constant dense<1> : tensor<i64>
    %c_162 = stablehlo.constant dense<0> : tensor<i64>
    %1528 = stablehlo.broadcast_in_dim %c_162, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %1529 = stablehlo.compare  LT, %1527, %1528,  SIGNED : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1>
    %c_163 = stablehlo.constant dense<3> : tensor<i64>
    %1530 = stablehlo.broadcast_in_dim %c_163, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %1531 = stablehlo.add %1527, %1530 : tensor<4xi64>
    %1532 = stablehlo.select %1529, %1531, %1527 : tensor<4xi1>, tensor<4xi64>
    %1533 = stablehlo.convert %1532 : (tensor<4xi64>) -> tensor<4xi32>
    %1534 = stablehlo.broadcast_in_dim %1533, dims = [0] : (tensor<4xi32>) -> tensor<4x1xi32>
    %1535 = stablehlo.broadcast_in_dim %c_161, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %1536 = "stablehlo.scatter"(%1526, %1534, %1535) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<3xi64>, tensor<4x1xi32>, tensor<4xi64>) -> tensor<3xi64>
    %1537 = call @_cumulative_reduction_9(%1536) : (tensor<3xi64>) -> tensor<3xi64>
    %c_164 = stablehlo.constant dense<1> : tensor<i64>
    %1538 = call @floor_divide(%1537, %c_164) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_165 = stablehlo.constant dense<4> : tensor<i64>
    %1539 = call @remainder_12(%1538, %c_165) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_166 = stablehlo.constant dense<0> : tensor<i64>
    %1540 = stablehlo.broadcast_in_dim %c_166, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1541 = stablehlo.compare  LT, %1514, %1540,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_167 = stablehlo.constant dense<4> : tensor<i64>
    %1542 = stablehlo.broadcast_in_dim %c_167, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1543 = stablehlo.add %1514, %1542 : tensor<1xi64>
    %1544 = stablehlo.select %1541, %1543, %1514 : tensor<1xi1>, tensor<1xi64>
    %1545 = stablehlo.convert %1544 : (tensor<1xi64>) -> tensor<1xi32>
    %1546 = stablehlo.broadcast_in_dim %1545, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %1547 = stablehlo.broadcast_in_dim %1496, dims = [1] : (tensor<4xf64>) -> tensor<1x4xf64>
    %1548 = "stablehlo.scatter"(%1547, %1546, %1491) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<1x1xi32>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %c_168 = stablehlo.constant dense<0> : tensor<i64>
    %1549 = stablehlo.broadcast_in_dim %c_168, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %1550 = stablehlo.compare  LT, %1539, %1549,  SIGNED : (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi1>
    %c_169 = stablehlo.constant dense<4> : tensor<i64>
    %1551 = stablehlo.broadcast_in_dim %c_169, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %1552 = stablehlo.add %1539, %1551 : tensor<3xi64>
    %1553 = stablehlo.select %1550, %1552, %1539 : tensor<3xi1>, tensor<3xi64>
    %1554 = stablehlo.convert %1553 : (tensor<3xi64>) -> tensor<3xi32>
    %1555 = stablehlo.broadcast_in_dim %1554, dims = [0] : (tensor<3xi32>) -> tensor<3x1xi32>
    %1556 = "stablehlo.scatter"(%1548, %1555, %1490) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<3x1xi32>, tensor<1x3xf64>) -> tensor<1x4xf64>
    %1557 = stablehlo.slice %1423 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1558 = stablehlo.reshape %1557 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1559 = stablehlo.slice %1556 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1560 = stablehlo.reshape %1559 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1561 = stablehlo.multiply %1558, %1560 : tensor<1xf64>
    %1562 = stablehlo.slice %1423 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1563 = stablehlo.reshape %1562 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1564 = stablehlo.slice %1556 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1565 = stablehlo.reshape %1564 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1566 = stablehlo.multiply %1563, %1565 : tensor<1xf64>
    %1567 = stablehlo.subtract %1561, %1566 : tensor<1xf64>
    %1568 = stablehlo.slice %1423 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1569 = stablehlo.reshape %1568 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1570 = stablehlo.slice %1556 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1571 = stablehlo.reshape %1570 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1572 = stablehlo.multiply %1569, %1571 : tensor<1xf64>
    %1573 = stablehlo.subtract %1567, %1572 : tensor<1xf64>
    %1574 = stablehlo.slice %1423 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1575 = stablehlo.reshape %1574 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1576 = stablehlo.slice %1556 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1577 = stablehlo.reshape %1576 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1578 = stablehlo.multiply %1575, %1577 : tensor<1xf64>
    %1579 = stablehlo.subtract %1573, %1578 : tensor<1xf64>
    %1580 = stablehlo.slice %1423 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1581 = stablehlo.reshape %1580 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1582 = stablehlo.slice %1556 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1583 = stablehlo.reshape %1582 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1584 = stablehlo.multiply %1581, %1583 : tensor<1xf64>
    %1585 = stablehlo.slice %1423 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1586 = stablehlo.reshape %1585 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1587 = stablehlo.slice %1556 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1588 = stablehlo.reshape %1587 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1589 = stablehlo.multiply %1586, %1588 : tensor<1xf64>
    %1590 = stablehlo.add %1584, %1589 : tensor<1xf64>
    %1591 = stablehlo.slice %1423 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1592 = stablehlo.reshape %1591 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1593 = stablehlo.slice %1556 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1594 = stablehlo.reshape %1593 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1595 = stablehlo.multiply %1592, %1594 : tensor<1xf64>
    %1596 = stablehlo.add %1590, %1595 : tensor<1xf64>
    %1597 = stablehlo.slice %1423 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1598 = stablehlo.reshape %1597 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1599 = stablehlo.slice %1556 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1600 = stablehlo.reshape %1599 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1601 = stablehlo.multiply %1598, %1600 : tensor<1xf64>
    %1602 = stablehlo.subtract %1596, %1601 : tensor<1xf64>
    %1603 = stablehlo.slice %1423 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1604 = stablehlo.reshape %1603 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1605 = stablehlo.slice %1556 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1606 = stablehlo.reshape %1605 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1607 = stablehlo.multiply %1604, %1606 : tensor<1xf64>
    %1608 = stablehlo.slice %1423 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1609 = stablehlo.reshape %1608 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1610 = stablehlo.slice %1556 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1611 = stablehlo.reshape %1610 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1612 = stablehlo.multiply %1609, %1611 : tensor<1xf64>
    %1613 = stablehlo.subtract %1607, %1612 : tensor<1xf64>
    %1614 = stablehlo.slice %1423 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1615 = stablehlo.reshape %1614 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1616 = stablehlo.slice %1556 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1617 = stablehlo.reshape %1616 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1618 = stablehlo.multiply %1615, %1617 : tensor<1xf64>
    %1619 = stablehlo.add %1613, %1618 : tensor<1xf64>
    %1620 = stablehlo.slice %1423 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1621 = stablehlo.reshape %1620 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1622 = stablehlo.slice %1556 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1623 = stablehlo.reshape %1622 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1624 = stablehlo.multiply %1621, %1623 : tensor<1xf64>
    %1625 = stablehlo.add %1619, %1624 : tensor<1xf64>
    %1626 = stablehlo.slice %1423 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1627 = stablehlo.reshape %1626 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1628 = stablehlo.slice %1556 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1629 = stablehlo.reshape %1628 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1630 = stablehlo.multiply %1627, %1629 : tensor<1xf64>
    %1631 = stablehlo.slice %1423 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1632 = stablehlo.reshape %1631 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1633 = stablehlo.slice %1556 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1634 = stablehlo.reshape %1633 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1635 = stablehlo.multiply %1632, %1634 : tensor<1xf64>
    %1636 = stablehlo.add %1630, %1635 : tensor<1xf64>
    %1637 = stablehlo.slice %1423 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1638 = stablehlo.reshape %1637 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1639 = stablehlo.slice %1556 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1640 = stablehlo.reshape %1639 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1641 = stablehlo.multiply %1638, %1640 : tensor<1xf64>
    %1642 = stablehlo.subtract %1636, %1641 : tensor<1xf64>
    %1643 = stablehlo.slice %1423 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1644 = stablehlo.reshape %1643 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1645 = stablehlo.slice %1556 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1646 = stablehlo.reshape %1645 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1647 = stablehlo.multiply %1644, %1646 : tensor<1xf64>
    %1648 = stablehlo.add %1642, %1647 : tensor<1xf64>
    %1649 = stablehlo.broadcast_in_dim %1579, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1650 = stablehlo.broadcast_in_dim %1602, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1651 = stablehlo.broadcast_in_dim %1625, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1652 = stablehlo.broadcast_in_dim %1648, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1653 = stablehlo.concatenate %1649, %1650, %1651, %1652, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %1654 = stablehlo.slice %1297 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %1655 = stablehlo.reshape %1654 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %1656 = stablehlo.slice %1653 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1657 = stablehlo.reshape %1656 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1658 = stablehlo.slice %1653 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %1659 = stablehlo.dot_general %1658, %1655, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1660 = stablehlo.broadcast_in_dim %1659, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1661 = stablehlo.broadcast_in_dim %1660, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1662 = stablehlo.multiply %1661, %1658 : tensor<1x3xf64>
    %cst_170 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1663 = stablehlo.broadcast_in_dim %cst_170, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %1664 = stablehlo.multiply %1663, %1662 : tensor<1x3xf64>
    %1665 = stablehlo.multiply %1657, %1657 : tensor<1xf64>
    %1666 = stablehlo.dot_general %1658, %1658, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1667 = stablehlo.subtract %1665, %1666 : tensor<1xf64>
    %1668 = stablehlo.broadcast_in_dim %1667, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1669 = stablehlo.broadcast_in_dim %1668, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1670 = stablehlo.multiply %1669, %1655 : tensor<1x3xf64>
    %1671 = stablehlo.add %1664, %1670 : tensor<1x3xf64>
    %cst_171 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1672 = stablehlo.broadcast_in_dim %cst_171, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1673 = stablehlo.multiply %1672, %1657 : tensor<1xf64>
    %1674 = call @cross(%1658, %1655) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %1675 = stablehlo.broadcast_in_dim %1673, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1676 = stablehlo.broadcast_in_dim %1675, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1677 = stablehlo.multiply %1676, %1674 : tensor<1x3xf64>
    %1678 = stablehlo.add %1671, %1677 : tensor<1x3xf64>
    %1679 = stablehlo.subtract %1449, %1678 : tensor<1x3xf64>
    %1680 = stablehlo.broadcast_in_dim %1449, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %1681 = stablehlo.broadcast_in_dim %1474, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %1682 = stablehlo.broadcast_in_dim %1653, dims = [0, 1] : (tensor<1x4xf64>) -> tensor<1x4x1xf64>
    %1683 = stablehlo.broadcast_in_dim %1653, dims = [0, 2] : (tensor<1x4xf64>) -> tensor<1x1x4xf64>
    %1684 = stablehlo.broadcast_in_dim %1682, dims = [0, 1, 2] : (tensor<1x4x1xf64>) -> tensor<1x4x4xf64>
    %1685 = stablehlo.broadcast_in_dim %1683, dims = [0, 1, 2] : (tensor<1x1x4xf64>) -> tensor<1x4x4xf64>
    %1686 = stablehlo.multiply %1684, %1685 : tensor<1x4x4xf64>
    %1687 = stablehlo.slice %1686 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1688 = stablehlo.reshape %1687 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1689 = stablehlo.slice %1686 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1690 = stablehlo.reshape %1689 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1691 = stablehlo.add %1688, %1690 : tensor<1xf64>
    %1692 = stablehlo.slice %1686 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1693 = stablehlo.reshape %1692 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1694 = stablehlo.subtract %1691, %1693 : tensor<1xf64>
    %1695 = stablehlo.slice %1686 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1696 = stablehlo.reshape %1695 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1697 = stablehlo.subtract %1694, %1696 : tensor<1xf64>
    %1698 = stablehlo.slice %1686 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1699 = stablehlo.reshape %1698 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1700 = stablehlo.slice %1686 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1701 = stablehlo.reshape %1700 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1702 = stablehlo.subtract %1699, %1701 : tensor<1xf64>
    %cst_172 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1703 = stablehlo.broadcast_in_dim %cst_172, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1704 = stablehlo.multiply %1703, %1702 : tensor<1xf64>
    %1705 = stablehlo.slice %1686 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1706 = stablehlo.reshape %1705 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1707 = stablehlo.slice %1686 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1708 = stablehlo.reshape %1707 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1709 = stablehlo.add %1706, %1708 : tensor<1xf64>
    %cst_173 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1710 = stablehlo.broadcast_in_dim %cst_173, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1711 = stablehlo.multiply %1710, %1709 : tensor<1xf64>
    %1712 = stablehlo.slice %1686 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1713 = stablehlo.reshape %1712 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1714 = stablehlo.slice %1686 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1715 = stablehlo.reshape %1714 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1716 = stablehlo.add %1713, %1715 : tensor<1xf64>
    %cst_174 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1717 = stablehlo.broadcast_in_dim %cst_174, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1718 = stablehlo.multiply %1717, %1716 : tensor<1xf64>
    %1719 = stablehlo.slice %1686 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1720 = stablehlo.reshape %1719 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1721 = stablehlo.slice %1686 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1722 = stablehlo.reshape %1721 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1723 = stablehlo.subtract %1720, %1722 : tensor<1xf64>
    %1724 = stablehlo.slice %1686 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1725 = stablehlo.reshape %1724 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1726 = stablehlo.add %1723, %1725 : tensor<1xf64>
    %1727 = stablehlo.slice %1686 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1728 = stablehlo.reshape %1727 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1729 = stablehlo.subtract %1726, %1728 : tensor<1xf64>
    %1730 = stablehlo.slice %1686 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1731 = stablehlo.reshape %1730 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1732 = stablehlo.slice %1686 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1733 = stablehlo.reshape %1732 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1734 = stablehlo.subtract %1731, %1733 : tensor<1xf64>
    %cst_175 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1735 = stablehlo.broadcast_in_dim %cst_175, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1736 = stablehlo.multiply %1735, %1734 : tensor<1xf64>
    %1737 = stablehlo.slice %1686 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1738 = stablehlo.reshape %1737 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1739 = stablehlo.slice %1686 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1740 = stablehlo.reshape %1739 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1741 = stablehlo.subtract %1738, %1740 : tensor<1xf64>
    %cst_176 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1742 = stablehlo.broadcast_in_dim %cst_176, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1743 = stablehlo.multiply %1742, %1741 : tensor<1xf64>
    %1744 = stablehlo.slice %1686 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1745 = stablehlo.reshape %1744 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1746 = stablehlo.slice %1686 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1747 = stablehlo.reshape %1746 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1748 = stablehlo.add %1745, %1747 : tensor<1xf64>
    %cst_177 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1749 = stablehlo.broadcast_in_dim %cst_177, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1750 = stablehlo.multiply %1749, %1748 : tensor<1xf64>
    %1751 = stablehlo.slice %1686 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1752 = stablehlo.reshape %1751 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1753 = stablehlo.slice %1686 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1754 = stablehlo.reshape %1753 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1755 = stablehlo.subtract %1752, %1754 : tensor<1xf64>
    %1756 = stablehlo.slice %1686 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1757 = stablehlo.reshape %1756 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1758 = stablehlo.subtract %1755, %1757 : tensor<1xf64>
    %1759 = stablehlo.slice %1686 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %1760 = stablehlo.reshape %1759 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %1761 = stablehlo.add %1758, %1760 : tensor<1xf64>
    %1762 = stablehlo.broadcast_in_dim %1697, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1763 = stablehlo.broadcast_in_dim %1704, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1764 = stablehlo.broadcast_in_dim %1711, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1765 = stablehlo.concatenate %1762, %1763, %1764, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %1766 = stablehlo.broadcast_in_dim %1718, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1767 = stablehlo.broadcast_in_dim %1729, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1768 = stablehlo.broadcast_in_dim %1736, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1769 = stablehlo.concatenate %1766, %1767, %1768, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %1770 = stablehlo.broadcast_in_dim %1743, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1771 = stablehlo.broadcast_in_dim %1750, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1772 = stablehlo.broadcast_in_dim %1761, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1773 = stablehlo.concatenate %1770, %1771, %1772, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %1774 = stablehlo.broadcast_in_dim %1765, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %1775 = stablehlo.broadcast_in_dim %1769, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %1776 = stablehlo.broadcast_in_dim %1773, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %1777 = stablehlo.concatenate %1774, %1775, %1776, dim = 1 : (tensor<1x1x3xf64>, tensor<1x1x3xf64>, tensor<1x1x3xf64>) -> tensor<1x3x3xf64>
    %1778 = call @_take_2(%1679, %c_16) : (tensor<1x3xf64>, tensor<1xi64>) -> tensor<1x3xf64>
    %1779 = call @_take_3(%1653, %c_16) : (tensor<1x4xf64>, tensor<1xi64>) -> tensor<1x4xf64>
    %1780 = call @_take_4(%arg5, %c_17) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %1781 = call @_take_4(%arg6, %c_18) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %1782 = call @_take_6(%arg13, %c_19) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %1783 = call @_take_6(%arg0, %c_20) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %1784 = stablehlo.slice %arg1 [5:6, 0:3] : (tensor<9x3xf64>) -> tensor<1x3xf64>
    %1785 = stablehlo.slice %arg2 [5:6, 0:4] : (tensor<9x4xf64>) -> tensor<1x4xf64>
    %1786 = stablehlo.slice %1779 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1787 = stablehlo.reshape %1786 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1788 = stablehlo.slice %1779 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %1789 = stablehlo.dot_general %1788, %1784, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1790 = stablehlo.broadcast_in_dim %1789, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1791 = stablehlo.broadcast_in_dim %1790, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1792 = stablehlo.multiply %1791, %1788 : tensor<1x3xf64>
    %cst_178 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1793 = stablehlo.broadcast_in_dim %cst_178, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %1794 = stablehlo.multiply %1793, %1792 : tensor<1x3xf64>
    %1795 = stablehlo.multiply %1787, %1787 : tensor<1xf64>
    %1796 = stablehlo.dot_general %1788, %1788, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1797 = stablehlo.subtract %1795, %1796 : tensor<1xf64>
    %1798 = stablehlo.broadcast_in_dim %1797, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1799 = stablehlo.broadcast_in_dim %1798, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1800 = stablehlo.multiply %1799, %1784 : tensor<1x3xf64>
    %1801 = stablehlo.add %1794, %1800 : tensor<1x3xf64>
    %cst_179 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1802 = stablehlo.broadcast_in_dim %cst_179, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1803 = stablehlo.multiply %1802, %1787 : tensor<1xf64>
    %1804 = call @cross(%1788, %1784) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %1805 = stablehlo.broadcast_in_dim %1803, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1806 = stablehlo.broadcast_in_dim %1805, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1807 = stablehlo.multiply %1806, %1804 : tensor<1x3xf64>
    %1808 = stablehlo.add %1801, %1807 : tensor<1x3xf64>
    %1809 = stablehlo.add %1778, %1808 : tensor<1x3xf64>
    %1810 = stablehlo.slice %1779 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1811 = stablehlo.reshape %1810 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1812 = stablehlo.slice %1785 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1813 = stablehlo.reshape %1812 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1814 = stablehlo.multiply %1811, %1813 : tensor<1xf64>
    %1815 = stablehlo.slice %1779 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1816 = stablehlo.reshape %1815 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1817 = stablehlo.slice %1785 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1818 = stablehlo.reshape %1817 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1819 = stablehlo.multiply %1816, %1818 : tensor<1xf64>
    %1820 = stablehlo.subtract %1814, %1819 : tensor<1xf64>
    %1821 = stablehlo.slice %1779 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1822 = stablehlo.reshape %1821 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1823 = stablehlo.slice %1785 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1824 = stablehlo.reshape %1823 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1825 = stablehlo.multiply %1822, %1824 : tensor<1xf64>
    %1826 = stablehlo.subtract %1820, %1825 : tensor<1xf64>
    %1827 = stablehlo.slice %1779 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1828 = stablehlo.reshape %1827 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1829 = stablehlo.slice %1785 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1830 = stablehlo.reshape %1829 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1831 = stablehlo.multiply %1828, %1830 : tensor<1xf64>
    %1832 = stablehlo.subtract %1826, %1831 : tensor<1xf64>
    %1833 = stablehlo.slice %1779 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1834 = stablehlo.reshape %1833 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1835 = stablehlo.slice %1785 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1836 = stablehlo.reshape %1835 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1837 = stablehlo.multiply %1834, %1836 : tensor<1xf64>
    %1838 = stablehlo.slice %1779 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1839 = stablehlo.reshape %1838 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1840 = stablehlo.slice %1785 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1841 = stablehlo.reshape %1840 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1842 = stablehlo.multiply %1839, %1841 : tensor<1xf64>
    %1843 = stablehlo.add %1837, %1842 : tensor<1xf64>
    %1844 = stablehlo.slice %1779 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1845 = stablehlo.reshape %1844 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1846 = stablehlo.slice %1785 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1847 = stablehlo.reshape %1846 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1848 = stablehlo.multiply %1845, %1847 : tensor<1xf64>
    %1849 = stablehlo.add %1843, %1848 : tensor<1xf64>
    %1850 = stablehlo.slice %1779 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1851 = stablehlo.reshape %1850 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1852 = stablehlo.slice %1785 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1853 = stablehlo.reshape %1852 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1854 = stablehlo.multiply %1851, %1853 : tensor<1xf64>
    %1855 = stablehlo.subtract %1849, %1854 : tensor<1xf64>
    %1856 = stablehlo.slice %1779 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1857 = stablehlo.reshape %1856 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1858 = stablehlo.slice %1785 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1859 = stablehlo.reshape %1858 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1860 = stablehlo.multiply %1857, %1859 : tensor<1xf64>
    %1861 = stablehlo.slice %1779 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1862 = stablehlo.reshape %1861 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1863 = stablehlo.slice %1785 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1864 = stablehlo.reshape %1863 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1865 = stablehlo.multiply %1862, %1864 : tensor<1xf64>
    %1866 = stablehlo.subtract %1860, %1865 : tensor<1xf64>
    %1867 = stablehlo.slice %1779 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1868 = stablehlo.reshape %1867 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1869 = stablehlo.slice %1785 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1870 = stablehlo.reshape %1869 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1871 = stablehlo.multiply %1868, %1870 : tensor<1xf64>
    %1872 = stablehlo.add %1866, %1871 : tensor<1xf64>
    %1873 = stablehlo.slice %1779 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1874 = stablehlo.reshape %1873 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1875 = stablehlo.slice %1785 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1876 = stablehlo.reshape %1875 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1877 = stablehlo.multiply %1874, %1876 : tensor<1xf64>
    %1878 = stablehlo.add %1872, %1877 : tensor<1xf64>
    %1879 = stablehlo.slice %1779 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1880 = stablehlo.reshape %1879 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1881 = stablehlo.slice %1785 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1882 = stablehlo.reshape %1881 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1883 = stablehlo.multiply %1880, %1882 : tensor<1xf64>
    %1884 = stablehlo.slice %1779 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1885 = stablehlo.reshape %1884 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1886 = stablehlo.slice %1785 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1887 = stablehlo.reshape %1886 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1888 = stablehlo.multiply %1885, %1887 : tensor<1xf64>
    %1889 = stablehlo.add %1883, %1888 : tensor<1xf64>
    %1890 = stablehlo.slice %1779 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1891 = stablehlo.reshape %1890 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1892 = stablehlo.slice %1785 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1893 = stablehlo.reshape %1892 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1894 = stablehlo.multiply %1891, %1893 : tensor<1xf64>
    %1895 = stablehlo.subtract %1889, %1894 : tensor<1xf64>
    %1896 = stablehlo.slice %1779 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1897 = stablehlo.reshape %1896 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1898 = stablehlo.slice %1785 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1899 = stablehlo.reshape %1898 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1900 = stablehlo.multiply %1897, %1899 : tensor<1xf64>
    %1901 = stablehlo.add %1895, %1900 : tensor<1xf64>
    %1902 = stablehlo.broadcast_in_dim %1832, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1903 = stablehlo.broadcast_in_dim %1855, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1904 = stablehlo.broadcast_in_dim %1878, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1905 = stablehlo.broadcast_in_dim %1901, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1906 = stablehlo.concatenate %1902, %1903, %1904, %1905, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %1907 = stablehlo.slice %1780 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %1908 = stablehlo.reshape %1907 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %1909 = stablehlo.slice %1906 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1910 = stablehlo.reshape %1909 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1911 = stablehlo.slice %1906 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %1912 = stablehlo.dot_general %1911, %1908, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1913 = stablehlo.broadcast_in_dim %1912, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1914 = stablehlo.broadcast_in_dim %1913, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1915 = stablehlo.multiply %1914, %1911 : tensor<1x3xf64>
    %cst_180 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1916 = stablehlo.broadcast_in_dim %cst_180, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %1917 = stablehlo.multiply %1916, %1915 : tensor<1x3xf64>
    %1918 = stablehlo.multiply %1910, %1910 : tensor<1xf64>
    %1919 = stablehlo.dot_general %1911, %1911, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1920 = stablehlo.subtract %1918, %1919 : tensor<1xf64>
    %1921 = stablehlo.broadcast_in_dim %1920, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1922 = stablehlo.broadcast_in_dim %1921, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1923 = stablehlo.multiply %1922, %1908 : tensor<1x3xf64>
    %1924 = stablehlo.add %1917, %1923 : tensor<1x3xf64>
    %cst_181 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1925 = stablehlo.broadcast_in_dim %cst_181, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1926 = stablehlo.multiply %1925, %1910 : tensor<1xf64>
    %1927 = call @cross(%1911, %1908) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %1928 = stablehlo.broadcast_in_dim %1926, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1929 = stablehlo.broadcast_in_dim %1928, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1930 = stablehlo.multiply %1929, %1927 : tensor<1x3xf64>
    %1931 = stablehlo.add %1924, %1930 : tensor<1x3xf64>
    %1932 = stablehlo.add %1931, %1809 : tensor<1x3xf64>
    %1933 = stablehlo.slice %1781 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %1934 = stablehlo.reshape %1933 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %1935 = stablehlo.slice %1906 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %1936 = stablehlo.reshape %1935 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1937 = stablehlo.slice %1906 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %1938 = stablehlo.dot_general %1937, %1934, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1939 = stablehlo.broadcast_in_dim %1938, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1940 = stablehlo.broadcast_in_dim %1939, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1941 = stablehlo.multiply %1940, %1937 : tensor<1x3xf64>
    %cst_182 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1942 = stablehlo.broadcast_in_dim %cst_182, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %1943 = stablehlo.multiply %1942, %1941 : tensor<1x3xf64>
    %1944 = stablehlo.multiply %1936, %1936 : tensor<1xf64>
    %1945 = stablehlo.dot_general %1937, %1937, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %1946 = stablehlo.subtract %1944, %1945 : tensor<1xf64>
    %1947 = stablehlo.broadcast_in_dim %1946, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1948 = stablehlo.broadcast_in_dim %1947, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1949 = stablehlo.multiply %1948, %1934 : tensor<1x3xf64>
    %1950 = stablehlo.add %1943, %1949 : tensor<1x3xf64>
    %cst_183 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %1951 = stablehlo.broadcast_in_dim %cst_183, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1952 = stablehlo.multiply %1951, %1936 : tensor<1xf64>
    %1953 = call @cross(%1937, %1934) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %1954 = stablehlo.broadcast_in_dim %1952, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1955 = stablehlo.broadcast_in_dim %1954, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1956 = stablehlo.multiply %1955, %1953 : tensor<1x3xf64>
    %1957 = stablehlo.add %1950, %1956 : tensor<1x3xf64>
    %1958 = stablehlo.slice %1782 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %1959 = stablehlo.reshape %1958 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1960 = stablehlo.slice %1783 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %1961 = stablehlo.reshape %1960 : (tensor<1x1xf64>) -> tensor<1xf64>
    %1962 = stablehlo.subtract %1959, %1961 : tensor<1xf64>
    %1963 = stablehlo.slice %1781 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %1964 = stablehlo.reshape %1963 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %cst_184 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %1965 = stablehlo.broadcast_in_dim %cst_184, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1966 = stablehlo.multiply %1962, %1965 : tensor<1xf64>
    %1967 = stablehlo.sine %1966 : tensor<1xf64>
    %cst_185 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %1968 = stablehlo.broadcast_in_dim %cst_185, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %1969 = stablehlo.multiply %1962, %1968 : tensor<1xf64>
    %1970 = stablehlo.cosine %1969 : tensor<1xf64>
    %1971 = stablehlo.broadcast_in_dim %1967, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %1972 = stablehlo.broadcast_in_dim %1971, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %1973 = stablehlo.multiply %1964, %1972 : tensor<1x3xf64>
    %1974 = stablehlo.broadcast_in_dim %1970, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %c_186 = stablehlo.constant dense<0> : tensor<i64>
    %1975 = stablehlo.reshape %c_186 : (tensor<i64>) -> tensor<1xi64>
    %1976 = stablehlo.slice %1975 [0:1] : (tensor<1xi64>) -> tensor<1xi64>
    %1977 = stablehlo.reshape %1976 : (tensor<1xi64>) -> tensor<i64>
    %1978 = stablehlo.broadcast_in_dim %1977, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %cst_187 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %1979 = stablehlo.broadcast_in_dim %cst_187, dims = [] : (tensor<f64>) -> tensor<4xf64>
    %c_188 = stablehlo.constant dense<0> : tensor<i64>
    %1980 = stablehlo.broadcast_in_dim %c_188, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1981 = stablehlo.compare  LT, %1978, %1980,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_189 = stablehlo.constant dense<3> : tensor<i64>
    %1982 = stablehlo.broadcast_in_dim %c_189, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1983 = stablehlo.add %1978, %1982 : tensor<1xi64>
    %1984 = call @_where_7(%1981, %1983, %1978) : (tensor<1xi1>, tensor<1xi64>, tensor<1xi64>) -> tensor<1xi64>
    %c_190 = stablehlo.constant dense<0> : tensor<i64>
    %c_191 = stablehlo.constant dense<3> : tensor<i64>
    %1985 = call @clip(%1984, %c_190, %c_191) : (tensor<1xi64>, tensor<i64>, tensor<i64>) -> tensor<1xi64>
    %1986 = call @argsort(%1985) : (tensor<1xi64>) -> tensor<1xi64>
    %1987 = stablehlo.iota dim = 0 : tensor<1xi64>
    %c_192 = stablehlo.constant dense<0> : tensor<i64>
    %1988 = stablehlo.broadcast_in_dim %c_192, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1989 = stablehlo.compare  LT, %1986, %1988,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_193 = stablehlo.constant dense<1> : tensor<i64>
    %1990 = stablehlo.broadcast_in_dim %c_193, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1991 = stablehlo.add %1986, %1990 : tensor<1xi64>
    %1992 = stablehlo.select %1989, %1991, %1986 : tensor<1xi1>, tensor<1xi64>
    %1993 = stablehlo.convert %1992 : (tensor<1xi64>) -> tensor<1xi32>
    %1994 = stablehlo.broadcast_in_dim %1993, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %1995 = stablehlo.convert %1985 : tensor<1xi64>
    %1996 = "stablehlo.scatter"(%1995, %1994, %1987) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<1xi64>, tensor<1x1xi32>, tensor<1xi64>) -> tensor<1xi64>
    %1997 = stablehlo.convert %1996 : tensor<1xi64>
    %c_194 = stablehlo.constant dense<true> : tensor<i1>
    %1998 = stablehlo.broadcast_in_dim %c_194, dims = [] : (tensor<i1>) -> tensor<4xi1>
    %c_195 = stablehlo.constant dense<0> : tensor<i64>
    %1999 = stablehlo.broadcast_in_dim %c_195, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2000 = stablehlo.compare  LT, %1997, %1999,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_196 = stablehlo.constant dense<4> : tensor<i64>
    %2001 = stablehlo.broadcast_in_dim %c_196, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2002 = stablehlo.add %1997, %2001 : tensor<1xi64>
    %2003 = stablehlo.select %2000, %2002, %1997 : tensor<1xi1>, tensor<1xi64>
    %2004 = stablehlo.convert %2003 : (tensor<1xi64>) -> tensor<1xi32>
    %2005 = stablehlo.broadcast_in_dim %2004, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %c_197 = stablehlo.constant dense<false> : tensor<i1>
    %2006 = stablehlo.broadcast_in_dim %c_197, dims = [] : (tensor<i1>) -> tensor<1xi1>
    %2007 = "stablehlo.scatter"(%1998, %2005, %2006) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i1>, %arg114: tensor<i1>):
      stablehlo.return %arg114 : tensor<i1>
    }) : (tensor<4xi1>, tensor<1x1xi32>, tensor<1xi1>) -> tensor<4xi1>
    %2008 = call @_cumulative_reduction(%2007) : (tensor<4xi1>) -> tensor<4xi64>
    %c_198 = stablehlo.constant dense<0> : tensor<i64>
    %2009 = stablehlo.broadcast_in_dim %c_198, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %c_199 = stablehlo.constant dense<0> : tensor<i64>
    %2010 = call @clip_8(%2008, %c_199) : (tensor<4xi64>, tensor<i64>) -> tensor<4xi64>
    %c_200 = stablehlo.constant dense<1> : tensor<i64>
    %c_201 = stablehlo.constant dense<0> : tensor<i64>
    %2011 = stablehlo.broadcast_in_dim %c_201, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %2012 = stablehlo.compare  LT, %2010, %2011,  SIGNED : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1>
    %c_202 = stablehlo.constant dense<3> : tensor<i64>
    %2013 = stablehlo.broadcast_in_dim %c_202, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %2014 = stablehlo.add %2010, %2013 : tensor<4xi64>
    %2015 = stablehlo.select %2012, %2014, %2010 : tensor<4xi1>, tensor<4xi64>
    %2016 = stablehlo.convert %2015 : (tensor<4xi64>) -> tensor<4xi32>
    %2017 = stablehlo.broadcast_in_dim %2016, dims = [0] : (tensor<4xi32>) -> tensor<4x1xi32>
    %2018 = stablehlo.broadcast_in_dim %c_200, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %2019 = "stablehlo.scatter"(%2009, %2017, %2018) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<3xi64>, tensor<4x1xi32>, tensor<4xi64>) -> tensor<3xi64>
    %2020 = call @_cumulative_reduction_9(%2019) : (tensor<3xi64>) -> tensor<3xi64>
    %c_203 = stablehlo.constant dense<1> : tensor<i64>
    %2021 = call @floor_divide(%2020, %c_203) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_204 = stablehlo.constant dense<4> : tensor<i64>
    %2022 = call @remainder_12(%2021, %c_204) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_205 = stablehlo.constant dense<0> : tensor<i64>
    %2023 = stablehlo.broadcast_in_dim %c_205, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2024 = stablehlo.compare  LT, %1997, %2023,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_206 = stablehlo.constant dense<4> : tensor<i64>
    %2025 = stablehlo.broadcast_in_dim %c_206, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2026 = stablehlo.add %1997, %2025 : tensor<1xi64>
    %2027 = stablehlo.select %2024, %2026, %1997 : tensor<1xi1>, tensor<1xi64>
    %2028 = stablehlo.convert %2027 : (tensor<1xi64>) -> tensor<1xi32>
    %2029 = stablehlo.broadcast_in_dim %2028, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %2030 = stablehlo.broadcast_in_dim %1979, dims = [1] : (tensor<4xf64>) -> tensor<1x4xf64>
    %2031 = "stablehlo.scatter"(%2030, %2029, %1974) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<1x1xi32>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %c_207 = stablehlo.constant dense<0> : tensor<i64>
    %2032 = stablehlo.broadcast_in_dim %c_207, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %2033 = stablehlo.compare  LT, %2022, %2032,  SIGNED : (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi1>
    %c_208 = stablehlo.constant dense<4> : tensor<i64>
    %2034 = stablehlo.broadcast_in_dim %c_208, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %2035 = stablehlo.add %2022, %2034 : tensor<3xi64>
    %2036 = stablehlo.select %2033, %2035, %2022 : tensor<3xi1>, tensor<3xi64>
    %2037 = stablehlo.convert %2036 : (tensor<3xi64>) -> tensor<3xi32>
    %2038 = stablehlo.broadcast_in_dim %2037, dims = [0] : (tensor<3xi32>) -> tensor<3x1xi32>
    %2039 = "stablehlo.scatter"(%2031, %2038, %1973) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<3x1xi32>, tensor<1x3xf64>) -> tensor<1x4xf64>
    %2040 = stablehlo.slice %1906 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2041 = stablehlo.reshape %2040 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2042 = stablehlo.slice %2039 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2043 = stablehlo.reshape %2042 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2044 = stablehlo.multiply %2041, %2043 : tensor<1xf64>
    %2045 = stablehlo.slice %1906 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2046 = stablehlo.reshape %2045 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2047 = stablehlo.slice %2039 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2048 = stablehlo.reshape %2047 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2049 = stablehlo.multiply %2046, %2048 : tensor<1xf64>
    %2050 = stablehlo.subtract %2044, %2049 : tensor<1xf64>
    %2051 = stablehlo.slice %1906 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2052 = stablehlo.reshape %2051 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2053 = stablehlo.slice %2039 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2054 = stablehlo.reshape %2053 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2055 = stablehlo.multiply %2052, %2054 : tensor<1xf64>
    %2056 = stablehlo.subtract %2050, %2055 : tensor<1xf64>
    %2057 = stablehlo.slice %1906 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2058 = stablehlo.reshape %2057 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2059 = stablehlo.slice %2039 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2060 = stablehlo.reshape %2059 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2061 = stablehlo.multiply %2058, %2060 : tensor<1xf64>
    %2062 = stablehlo.subtract %2056, %2061 : tensor<1xf64>
    %2063 = stablehlo.slice %1906 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2064 = stablehlo.reshape %2063 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2065 = stablehlo.slice %2039 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2066 = stablehlo.reshape %2065 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2067 = stablehlo.multiply %2064, %2066 : tensor<1xf64>
    %2068 = stablehlo.slice %1906 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2069 = stablehlo.reshape %2068 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2070 = stablehlo.slice %2039 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2071 = stablehlo.reshape %2070 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2072 = stablehlo.multiply %2069, %2071 : tensor<1xf64>
    %2073 = stablehlo.add %2067, %2072 : tensor<1xf64>
    %2074 = stablehlo.slice %1906 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2075 = stablehlo.reshape %2074 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2076 = stablehlo.slice %2039 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2077 = stablehlo.reshape %2076 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2078 = stablehlo.multiply %2075, %2077 : tensor<1xf64>
    %2079 = stablehlo.add %2073, %2078 : tensor<1xf64>
    %2080 = stablehlo.slice %1906 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2081 = stablehlo.reshape %2080 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2082 = stablehlo.slice %2039 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2083 = stablehlo.reshape %2082 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2084 = stablehlo.multiply %2081, %2083 : tensor<1xf64>
    %2085 = stablehlo.subtract %2079, %2084 : tensor<1xf64>
    %2086 = stablehlo.slice %1906 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2087 = stablehlo.reshape %2086 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2088 = stablehlo.slice %2039 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2089 = stablehlo.reshape %2088 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2090 = stablehlo.multiply %2087, %2089 : tensor<1xf64>
    %2091 = stablehlo.slice %1906 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2092 = stablehlo.reshape %2091 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2093 = stablehlo.slice %2039 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2094 = stablehlo.reshape %2093 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2095 = stablehlo.multiply %2092, %2094 : tensor<1xf64>
    %2096 = stablehlo.subtract %2090, %2095 : tensor<1xf64>
    %2097 = stablehlo.slice %1906 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2098 = stablehlo.reshape %2097 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2099 = stablehlo.slice %2039 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2100 = stablehlo.reshape %2099 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2101 = stablehlo.multiply %2098, %2100 : tensor<1xf64>
    %2102 = stablehlo.add %2096, %2101 : tensor<1xf64>
    %2103 = stablehlo.slice %1906 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2104 = stablehlo.reshape %2103 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2105 = stablehlo.slice %2039 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2106 = stablehlo.reshape %2105 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2107 = stablehlo.multiply %2104, %2106 : tensor<1xf64>
    %2108 = stablehlo.add %2102, %2107 : tensor<1xf64>
    %2109 = stablehlo.slice %1906 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2110 = stablehlo.reshape %2109 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2111 = stablehlo.slice %2039 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2112 = stablehlo.reshape %2111 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2113 = stablehlo.multiply %2110, %2112 : tensor<1xf64>
    %2114 = stablehlo.slice %1906 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2115 = stablehlo.reshape %2114 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2116 = stablehlo.slice %2039 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2117 = stablehlo.reshape %2116 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2118 = stablehlo.multiply %2115, %2117 : tensor<1xf64>
    %2119 = stablehlo.add %2113, %2118 : tensor<1xf64>
    %2120 = stablehlo.slice %1906 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2121 = stablehlo.reshape %2120 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2122 = stablehlo.slice %2039 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2123 = stablehlo.reshape %2122 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2124 = stablehlo.multiply %2121, %2123 : tensor<1xf64>
    %2125 = stablehlo.subtract %2119, %2124 : tensor<1xf64>
    %2126 = stablehlo.slice %1906 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2127 = stablehlo.reshape %2126 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2128 = stablehlo.slice %2039 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2129 = stablehlo.reshape %2128 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2130 = stablehlo.multiply %2127, %2129 : tensor<1xf64>
    %2131 = stablehlo.add %2125, %2130 : tensor<1xf64>
    %2132 = stablehlo.broadcast_in_dim %2062, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2133 = stablehlo.broadcast_in_dim %2085, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2134 = stablehlo.broadcast_in_dim %2108, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2135 = stablehlo.broadcast_in_dim %2131, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2136 = stablehlo.concatenate %2132, %2133, %2134, %2135, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %2137 = stablehlo.slice %1780 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %2138 = stablehlo.reshape %2137 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %2139 = stablehlo.slice %2136 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2140 = stablehlo.reshape %2139 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2141 = stablehlo.slice %2136 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %2142 = stablehlo.dot_general %2141, %2138, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2143 = stablehlo.broadcast_in_dim %2142, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2144 = stablehlo.broadcast_in_dim %2143, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2145 = stablehlo.multiply %2144, %2141 : tensor<1x3xf64>
    %cst_209 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2146 = stablehlo.broadcast_in_dim %cst_209, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %2147 = stablehlo.multiply %2146, %2145 : tensor<1x3xf64>
    %2148 = stablehlo.multiply %2140, %2140 : tensor<1xf64>
    %2149 = stablehlo.dot_general %2141, %2141, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2150 = stablehlo.subtract %2148, %2149 : tensor<1xf64>
    %2151 = stablehlo.broadcast_in_dim %2150, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2152 = stablehlo.broadcast_in_dim %2151, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2153 = stablehlo.multiply %2152, %2138 : tensor<1x3xf64>
    %2154 = stablehlo.add %2147, %2153 : tensor<1x3xf64>
    %cst_210 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2155 = stablehlo.broadcast_in_dim %cst_210, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2156 = stablehlo.multiply %2155, %2140 : tensor<1xf64>
    %2157 = call @cross(%2141, %2138) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %2158 = stablehlo.broadcast_in_dim %2156, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2159 = stablehlo.broadcast_in_dim %2158, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2160 = stablehlo.multiply %2159, %2157 : tensor<1x3xf64>
    %2161 = stablehlo.add %2154, %2160 : tensor<1x3xf64>
    %2162 = stablehlo.subtract %1932, %2161 : tensor<1x3xf64>
    %2163 = stablehlo.broadcast_in_dim %1932, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %2164 = stablehlo.broadcast_in_dim %1957, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %2165 = stablehlo.broadcast_in_dim %2136, dims = [0, 1] : (tensor<1x4xf64>) -> tensor<1x4x1xf64>
    %2166 = stablehlo.broadcast_in_dim %2136, dims = [0, 2] : (tensor<1x4xf64>) -> tensor<1x1x4xf64>
    %2167 = stablehlo.broadcast_in_dim %2165, dims = [0, 1, 2] : (tensor<1x4x1xf64>) -> tensor<1x4x4xf64>
    %2168 = stablehlo.broadcast_in_dim %2166, dims = [0, 1, 2] : (tensor<1x1x4xf64>) -> tensor<1x4x4xf64>
    %2169 = stablehlo.multiply %2167, %2168 : tensor<1x4x4xf64>
    %2170 = stablehlo.slice %2169 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2171 = stablehlo.reshape %2170 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2172 = stablehlo.slice %2169 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2173 = stablehlo.reshape %2172 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2174 = stablehlo.add %2171, %2173 : tensor<1xf64>
    %2175 = stablehlo.slice %2169 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2176 = stablehlo.reshape %2175 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2177 = stablehlo.subtract %2174, %2176 : tensor<1xf64>
    %2178 = stablehlo.slice %2169 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2179 = stablehlo.reshape %2178 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2180 = stablehlo.subtract %2177, %2179 : tensor<1xf64>
    %2181 = stablehlo.slice %2169 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2182 = stablehlo.reshape %2181 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2183 = stablehlo.slice %2169 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2184 = stablehlo.reshape %2183 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2185 = stablehlo.subtract %2182, %2184 : tensor<1xf64>
    %cst_211 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2186 = stablehlo.broadcast_in_dim %cst_211, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2187 = stablehlo.multiply %2186, %2185 : tensor<1xf64>
    %2188 = stablehlo.slice %2169 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2189 = stablehlo.reshape %2188 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2190 = stablehlo.slice %2169 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2191 = stablehlo.reshape %2190 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2192 = stablehlo.add %2189, %2191 : tensor<1xf64>
    %cst_212 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2193 = stablehlo.broadcast_in_dim %cst_212, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2194 = stablehlo.multiply %2193, %2192 : tensor<1xf64>
    %2195 = stablehlo.slice %2169 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2196 = stablehlo.reshape %2195 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2197 = stablehlo.slice %2169 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2198 = stablehlo.reshape %2197 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2199 = stablehlo.add %2196, %2198 : tensor<1xf64>
    %cst_213 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2200 = stablehlo.broadcast_in_dim %cst_213, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2201 = stablehlo.multiply %2200, %2199 : tensor<1xf64>
    %2202 = stablehlo.slice %2169 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2203 = stablehlo.reshape %2202 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2204 = stablehlo.slice %2169 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2205 = stablehlo.reshape %2204 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2206 = stablehlo.subtract %2203, %2205 : tensor<1xf64>
    %2207 = stablehlo.slice %2169 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2208 = stablehlo.reshape %2207 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2209 = stablehlo.add %2206, %2208 : tensor<1xf64>
    %2210 = stablehlo.slice %2169 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2211 = stablehlo.reshape %2210 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2212 = stablehlo.subtract %2209, %2211 : tensor<1xf64>
    %2213 = stablehlo.slice %2169 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2214 = stablehlo.reshape %2213 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2215 = stablehlo.slice %2169 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2216 = stablehlo.reshape %2215 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2217 = stablehlo.subtract %2214, %2216 : tensor<1xf64>
    %cst_214 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2218 = stablehlo.broadcast_in_dim %cst_214, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2219 = stablehlo.multiply %2218, %2217 : tensor<1xf64>
    %2220 = stablehlo.slice %2169 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2221 = stablehlo.reshape %2220 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2222 = stablehlo.slice %2169 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2223 = stablehlo.reshape %2222 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2224 = stablehlo.subtract %2221, %2223 : tensor<1xf64>
    %cst_215 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2225 = stablehlo.broadcast_in_dim %cst_215, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2226 = stablehlo.multiply %2225, %2224 : tensor<1xf64>
    %2227 = stablehlo.slice %2169 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2228 = stablehlo.reshape %2227 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2229 = stablehlo.slice %2169 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2230 = stablehlo.reshape %2229 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2231 = stablehlo.add %2228, %2230 : tensor<1xf64>
    %cst_216 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2232 = stablehlo.broadcast_in_dim %cst_216, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2233 = stablehlo.multiply %2232, %2231 : tensor<1xf64>
    %2234 = stablehlo.slice %2169 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2235 = stablehlo.reshape %2234 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2236 = stablehlo.slice %2169 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2237 = stablehlo.reshape %2236 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2238 = stablehlo.subtract %2235, %2237 : tensor<1xf64>
    %2239 = stablehlo.slice %2169 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2240 = stablehlo.reshape %2239 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2241 = stablehlo.subtract %2238, %2240 : tensor<1xf64>
    %2242 = stablehlo.slice %2169 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2243 = stablehlo.reshape %2242 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2244 = stablehlo.add %2241, %2243 : tensor<1xf64>
    %2245 = stablehlo.broadcast_in_dim %2180, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2246 = stablehlo.broadcast_in_dim %2187, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2247 = stablehlo.broadcast_in_dim %2194, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2248 = stablehlo.concatenate %2245, %2246, %2247, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %2249 = stablehlo.broadcast_in_dim %2201, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2250 = stablehlo.broadcast_in_dim %2212, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2251 = stablehlo.broadcast_in_dim %2219, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2252 = stablehlo.concatenate %2249, %2250, %2251, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %2253 = stablehlo.broadcast_in_dim %2226, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2254 = stablehlo.broadcast_in_dim %2233, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2255 = stablehlo.broadcast_in_dim %2244, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2256 = stablehlo.concatenate %2253, %2254, %2255, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %2257 = stablehlo.broadcast_in_dim %2248, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %2258 = stablehlo.broadcast_in_dim %2252, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %2259 = stablehlo.broadcast_in_dim %2256, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %2260 = stablehlo.concatenate %2257, %2258, %2259, dim = 1 : (tensor<1x1x3xf64>, tensor<1x1x3xf64>, tensor<1x1x3xf64>) -> tensor<1x3x3xf64>
    %2261 = call @_take_2(%2162, %c_21) : (tensor<1x3xf64>, tensor<1xi64>) -> tensor<1x3xf64>
    %2262 = call @_take_3(%2136, %c_21) : (tensor<1x4xf64>, tensor<1xi64>) -> tensor<1x4xf64>
    %2263 = call @_take_4(%arg5, %c_22) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %2264 = call @_take_4(%arg6, %c_23) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %2265 = call @_take_6(%arg13, %c_24) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %2266 = call @_take_6(%arg0, %c_25) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %2267 = stablehlo.slice %arg1 [6:7, 0:3] : (tensor<9x3xf64>) -> tensor<1x3xf64>
    %2268 = stablehlo.slice %arg2 [6:7, 0:4] : (tensor<9x4xf64>) -> tensor<1x4xf64>
    %2269 = stablehlo.slice %2262 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2270 = stablehlo.reshape %2269 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2271 = stablehlo.slice %2262 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %2272 = stablehlo.dot_general %2271, %2267, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2273 = stablehlo.broadcast_in_dim %2272, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2274 = stablehlo.broadcast_in_dim %2273, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2275 = stablehlo.multiply %2274, %2271 : tensor<1x3xf64>
    %cst_217 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2276 = stablehlo.broadcast_in_dim %cst_217, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %2277 = stablehlo.multiply %2276, %2275 : tensor<1x3xf64>
    %2278 = stablehlo.multiply %2270, %2270 : tensor<1xf64>
    %2279 = stablehlo.dot_general %2271, %2271, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2280 = stablehlo.subtract %2278, %2279 : tensor<1xf64>
    %2281 = stablehlo.broadcast_in_dim %2280, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2282 = stablehlo.broadcast_in_dim %2281, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2283 = stablehlo.multiply %2282, %2267 : tensor<1x3xf64>
    %2284 = stablehlo.add %2277, %2283 : tensor<1x3xf64>
    %cst_218 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2285 = stablehlo.broadcast_in_dim %cst_218, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2286 = stablehlo.multiply %2285, %2270 : tensor<1xf64>
    %2287 = call @cross(%2271, %2267) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %2288 = stablehlo.broadcast_in_dim %2286, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2289 = stablehlo.broadcast_in_dim %2288, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2290 = stablehlo.multiply %2289, %2287 : tensor<1x3xf64>
    %2291 = stablehlo.add %2284, %2290 : tensor<1x3xf64>
    %2292 = stablehlo.add %2261, %2291 : tensor<1x3xf64>
    %2293 = stablehlo.slice %2262 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2294 = stablehlo.reshape %2293 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2295 = stablehlo.slice %2268 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2296 = stablehlo.reshape %2295 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2297 = stablehlo.multiply %2294, %2296 : tensor<1xf64>
    %2298 = stablehlo.slice %2262 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2299 = stablehlo.reshape %2298 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2300 = stablehlo.slice %2268 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2301 = stablehlo.reshape %2300 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2302 = stablehlo.multiply %2299, %2301 : tensor<1xf64>
    %2303 = stablehlo.subtract %2297, %2302 : tensor<1xf64>
    %2304 = stablehlo.slice %2262 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2305 = stablehlo.reshape %2304 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2306 = stablehlo.slice %2268 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2307 = stablehlo.reshape %2306 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2308 = stablehlo.multiply %2305, %2307 : tensor<1xf64>
    %2309 = stablehlo.subtract %2303, %2308 : tensor<1xf64>
    %2310 = stablehlo.slice %2262 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2311 = stablehlo.reshape %2310 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2312 = stablehlo.slice %2268 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2313 = stablehlo.reshape %2312 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2314 = stablehlo.multiply %2311, %2313 : tensor<1xf64>
    %2315 = stablehlo.subtract %2309, %2314 : tensor<1xf64>
    %2316 = stablehlo.slice %2262 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2317 = stablehlo.reshape %2316 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2318 = stablehlo.slice %2268 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2319 = stablehlo.reshape %2318 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2320 = stablehlo.multiply %2317, %2319 : tensor<1xf64>
    %2321 = stablehlo.slice %2262 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2322 = stablehlo.reshape %2321 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2323 = stablehlo.slice %2268 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2324 = stablehlo.reshape %2323 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2325 = stablehlo.multiply %2322, %2324 : tensor<1xf64>
    %2326 = stablehlo.add %2320, %2325 : tensor<1xf64>
    %2327 = stablehlo.slice %2262 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2328 = stablehlo.reshape %2327 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2329 = stablehlo.slice %2268 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2330 = stablehlo.reshape %2329 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2331 = stablehlo.multiply %2328, %2330 : tensor<1xf64>
    %2332 = stablehlo.add %2326, %2331 : tensor<1xf64>
    %2333 = stablehlo.slice %2262 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2334 = stablehlo.reshape %2333 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2335 = stablehlo.slice %2268 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2336 = stablehlo.reshape %2335 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2337 = stablehlo.multiply %2334, %2336 : tensor<1xf64>
    %2338 = stablehlo.subtract %2332, %2337 : tensor<1xf64>
    %2339 = stablehlo.slice %2262 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2340 = stablehlo.reshape %2339 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2341 = stablehlo.slice %2268 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2342 = stablehlo.reshape %2341 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2343 = stablehlo.multiply %2340, %2342 : tensor<1xf64>
    %2344 = stablehlo.slice %2262 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2345 = stablehlo.reshape %2344 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2346 = stablehlo.slice %2268 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2347 = stablehlo.reshape %2346 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2348 = stablehlo.multiply %2345, %2347 : tensor<1xf64>
    %2349 = stablehlo.subtract %2343, %2348 : tensor<1xf64>
    %2350 = stablehlo.slice %2262 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2351 = stablehlo.reshape %2350 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2352 = stablehlo.slice %2268 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2353 = stablehlo.reshape %2352 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2354 = stablehlo.multiply %2351, %2353 : tensor<1xf64>
    %2355 = stablehlo.add %2349, %2354 : tensor<1xf64>
    %2356 = stablehlo.slice %2262 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2357 = stablehlo.reshape %2356 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2358 = stablehlo.slice %2268 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2359 = stablehlo.reshape %2358 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2360 = stablehlo.multiply %2357, %2359 : tensor<1xf64>
    %2361 = stablehlo.add %2355, %2360 : tensor<1xf64>
    %2362 = stablehlo.slice %2262 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2363 = stablehlo.reshape %2362 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2364 = stablehlo.slice %2268 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2365 = stablehlo.reshape %2364 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2366 = stablehlo.multiply %2363, %2365 : tensor<1xf64>
    %2367 = stablehlo.slice %2262 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2368 = stablehlo.reshape %2367 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2369 = stablehlo.slice %2268 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2370 = stablehlo.reshape %2369 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2371 = stablehlo.multiply %2368, %2370 : tensor<1xf64>
    %2372 = stablehlo.add %2366, %2371 : tensor<1xf64>
    %2373 = stablehlo.slice %2262 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2374 = stablehlo.reshape %2373 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2375 = stablehlo.slice %2268 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2376 = stablehlo.reshape %2375 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2377 = stablehlo.multiply %2374, %2376 : tensor<1xf64>
    %2378 = stablehlo.subtract %2372, %2377 : tensor<1xf64>
    %2379 = stablehlo.slice %2262 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2380 = stablehlo.reshape %2379 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2381 = stablehlo.slice %2268 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2382 = stablehlo.reshape %2381 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2383 = stablehlo.multiply %2380, %2382 : tensor<1xf64>
    %2384 = stablehlo.add %2378, %2383 : tensor<1xf64>
    %2385 = stablehlo.broadcast_in_dim %2315, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2386 = stablehlo.broadcast_in_dim %2338, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2387 = stablehlo.broadcast_in_dim %2361, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2388 = stablehlo.broadcast_in_dim %2384, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2389 = stablehlo.concatenate %2385, %2386, %2387, %2388, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %2390 = stablehlo.slice %2263 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %2391 = stablehlo.reshape %2390 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %2392 = stablehlo.slice %2389 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2393 = stablehlo.reshape %2392 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2394 = stablehlo.slice %2389 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %2395 = stablehlo.dot_general %2394, %2391, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2396 = stablehlo.broadcast_in_dim %2395, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2397 = stablehlo.broadcast_in_dim %2396, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2398 = stablehlo.multiply %2397, %2394 : tensor<1x3xf64>
    %cst_219 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2399 = stablehlo.broadcast_in_dim %cst_219, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %2400 = stablehlo.multiply %2399, %2398 : tensor<1x3xf64>
    %2401 = stablehlo.multiply %2393, %2393 : tensor<1xf64>
    %2402 = stablehlo.dot_general %2394, %2394, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2403 = stablehlo.subtract %2401, %2402 : tensor<1xf64>
    %2404 = stablehlo.broadcast_in_dim %2403, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2405 = stablehlo.broadcast_in_dim %2404, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2406 = stablehlo.multiply %2405, %2391 : tensor<1x3xf64>
    %2407 = stablehlo.add %2400, %2406 : tensor<1x3xf64>
    %cst_220 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2408 = stablehlo.broadcast_in_dim %cst_220, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2409 = stablehlo.multiply %2408, %2393 : tensor<1xf64>
    %2410 = call @cross(%2394, %2391) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %2411 = stablehlo.broadcast_in_dim %2409, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2412 = stablehlo.broadcast_in_dim %2411, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2413 = stablehlo.multiply %2412, %2410 : tensor<1x3xf64>
    %2414 = stablehlo.add %2407, %2413 : tensor<1x3xf64>
    %2415 = stablehlo.add %2414, %2292 : tensor<1x3xf64>
    %2416 = stablehlo.slice %2264 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %2417 = stablehlo.reshape %2416 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %2418 = stablehlo.slice %2389 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2419 = stablehlo.reshape %2418 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2420 = stablehlo.slice %2389 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %2421 = stablehlo.dot_general %2420, %2417, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2422 = stablehlo.broadcast_in_dim %2421, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2423 = stablehlo.broadcast_in_dim %2422, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2424 = stablehlo.multiply %2423, %2420 : tensor<1x3xf64>
    %cst_221 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2425 = stablehlo.broadcast_in_dim %cst_221, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %2426 = stablehlo.multiply %2425, %2424 : tensor<1x3xf64>
    %2427 = stablehlo.multiply %2419, %2419 : tensor<1xf64>
    %2428 = stablehlo.dot_general %2420, %2420, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2429 = stablehlo.subtract %2427, %2428 : tensor<1xf64>
    %2430 = stablehlo.broadcast_in_dim %2429, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2431 = stablehlo.broadcast_in_dim %2430, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2432 = stablehlo.multiply %2431, %2417 : tensor<1x3xf64>
    %2433 = stablehlo.add %2426, %2432 : tensor<1x3xf64>
    %cst_222 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2434 = stablehlo.broadcast_in_dim %cst_222, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2435 = stablehlo.multiply %2434, %2419 : tensor<1xf64>
    %2436 = call @cross(%2420, %2417) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %2437 = stablehlo.broadcast_in_dim %2435, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2438 = stablehlo.broadcast_in_dim %2437, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2439 = stablehlo.multiply %2438, %2436 : tensor<1x3xf64>
    %2440 = stablehlo.add %2433, %2439 : tensor<1x3xf64>
    %2441 = stablehlo.slice %2265 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %2442 = stablehlo.reshape %2441 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2443 = stablehlo.slice %2266 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %2444 = stablehlo.reshape %2443 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2445 = stablehlo.subtract %2442, %2444 : tensor<1xf64>
    %2446 = stablehlo.slice %2264 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %2447 = stablehlo.reshape %2446 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %cst_223 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %2448 = stablehlo.broadcast_in_dim %cst_223, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2449 = stablehlo.multiply %2445, %2448 : tensor<1xf64>
    %2450 = stablehlo.sine %2449 : tensor<1xf64>
    %cst_224 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %2451 = stablehlo.broadcast_in_dim %cst_224, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2452 = stablehlo.multiply %2445, %2451 : tensor<1xf64>
    %2453 = stablehlo.cosine %2452 : tensor<1xf64>
    %2454 = stablehlo.broadcast_in_dim %2450, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2455 = stablehlo.broadcast_in_dim %2454, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2456 = stablehlo.multiply %2447, %2455 : tensor<1x3xf64>
    %2457 = stablehlo.broadcast_in_dim %2453, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %c_225 = stablehlo.constant dense<0> : tensor<i64>
    %2458 = stablehlo.reshape %c_225 : (tensor<i64>) -> tensor<1xi64>
    %2459 = stablehlo.slice %2458 [0:1] : (tensor<1xi64>) -> tensor<1xi64>
    %2460 = stablehlo.reshape %2459 : (tensor<1xi64>) -> tensor<i64>
    %2461 = stablehlo.broadcast_in_dim %2460, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %cst_226 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %2462 = stablehlo.broadcast_in_dim %cst_226, dims = [] : (tensor<f64>) -> tensor<4xf64>
    %c_227 = stablehlo.constant dense<0> : tensor<i64>
    %2463 = stablehlo.broadcast_in_dim %c_227, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2464 = stablehlo.compare  LT, %2461, %2463,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_228 = stablehlo.constant dense<3> : tensor<i64>
    %2465 = stablehlo.broadcast_in_dim %c_228, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2466 = stablehlo.add %2461, %2465 : tensor<1xi64>
    %2467 = call @_where_7(%2464, %2466, %2461) : (tensor<1xi1>, tensor<1xi64>, tensor<1xi64>) -> tensor<1xi64>
    %c_229 = stablehlo.constant dense<0> : tensor<i64>
    %c_230 = stablehlo.constant dense<3> : tensor<i64>
    %2468 = call @clip(%2467, %c_229, %c_230) : (tensor<1xi64>, tensor<i64>, tensor<i64>) -> tensor<1xi64>
    %2469 = call @argsort(%2468) : (tensor<1xi64>) -> tensor<1xi64>
    %2470 = stablehlo.iota dim = 0 : tensor<1xi64>
    %c_231 = stablehlo.constant dense<0> : tensor<i64>
    %2471 = stablehlo.broadcast_in_dim %c_231, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2472 = stablehlo.compare  LT, %2469, %2471,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_232 = stablehlo.constant dense<1> : tensor<i64>
    %2473 = stablehlo.broadcast_in_dim %c_232, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2474 = stablehlo.add %2469, %2473 : tensor<1xi64>
    %2475 = stablehlo.select %2472, %2474, %2469 : tensor<1xi1>, tensor<1xi64>
    %2476 = stablehlo.convert %2475 : (tensor<1xi64>) -> tensor<1xi32>
    %2477 = stablehlo.broadcast_in_dim %2476, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %2478 = stablehlo.convert %2468 : tensor<1xi64>
    %2479 = "stablehlo.scatter"(%2478, %2477, %2470) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<1xi64>, tensor<1x1xi32>, tensor<1xi64>) -> tensor<1xi64>
    %2480 = stablehlo.convert %2479 : tensor<1xi64>
    %c_233 = stablehlo.constant dense<true> : tensor<i1>
    %2481 = stablehlo.broadcast_in_dim %c_233, dims = [] : (tensor<i1>) -> tensor<4xi1>
    %c_234 = stablehlo.constant dense<0> : tensor<i64>
    %2482 = stablehlo.broadcast_in_dim %c_234, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2483 = stablehlo.compare  LT, %2480, %2482,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_235 = stablehlo.constant dense<4> : tensor<i64>
    %2484 = stablehlo.broadcast_in_dim %c_235, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2485 = stablehlo.add %2480, %2484 : tensor<1xi64>
    %2486 = stablehlo.select %2483, %2485, %2480 : tensor<1xi1>, tensor<1xi64>
    %2487 = stablehlo.convert %2486 : (tensor<1xi64>) -> tensor<1xi32>
    %2488 = stablehlo.broadcast_in_dim %2487, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %c_236 = stablehlo.constant dense<false> : tensor<i1>
    %2489 = stablehlo.broadcast_in_dim %c_236, dims = [] : (tensor<i1>) -> tensor<1xi1>
    %2490 = "stablehlo.scatter"(%2481, %2488, %2489) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i1>, %arg114: tensor<i1>):
      stablehlo.return %arg114 : tensor<i1>
    }) : (tensor<4xi1>, tensor<1x1xi32>, tensor<1xi1>) -> tensor<4xi1>
    %2491 = call @_cumulative_reduction(%2490) : (tensor<4xi1>) -> tensor<4xi64>
    %c_237 = stablehlo.constant dense<0> : tensor<i64>
    %2492 = stablehlo.broadcast_in_dim %c_237, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %c_238 = stablehlo.constant dense<0> : tensor<i64>
    %2493 = call @clip_8(%2491, %c_238) : (tensor<4xi64>, tensor<i64>) -> tensor<4xi64>
    %c_239 = stablehlo.constant dense<1> : tensor<i64>
    %c_240 = stablehlo.constant dense<0> : tensor<i64>
    %2494 = stablehlo.broadcast_in_dim %c_240, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %2495 = stablehlo.compare  LT, %2493, %2494,  SIGNED : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1>
    %c_241 = stablehlo.constant dense<3> : tensor<i64>
    %2496 = stablehlo.broadcast_in_dim %c_241, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %2497 = stablehlo.add %2493, %2496 : tensor<4xi64>
    %2498 = stablehlo.select %2495, %2497, %2493 : tensor<4xi1>, tensor<4xi64>
    %2499 = stablehlo.convert %2498 : (tensor<4xi64>) -> tensor<4xi32>
    %2500 = stablehlo.broadcast_in_dim %2499, dims = [0] : (tensor<4xi32>) -> tensor<4x1xi32>
    %2501 = stablehlo.broadcast_in_dim %c_239, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %2502 = "stablehlo.scatter"(%2492, %2500, %2501) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<3xi64>, tensor<4x1xi32>, tensor<4xi64>) -> tensor<3xi64>
    %2503 = call @_cumulative_reduction_9(%2502) : (tensor<3xi64>) -> tensor<3xi64>
    %c_242 = stablehlo.constant dense<1> : tensor<i64>
    %2504 = call @floor_divide(%2503, %c_242) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_243 = stablehlo.constant dense<4> : tensor<i64>
    %2505 = call @remainder_12(%2504, %c_243) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_244 = stablehlo.constant dense<0> : tensor<i64>
    %2506 = stablehlo.broadcast_in_dim %c_244, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2507 = stablehlo.compare  LT, %2480, %2506,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_245 = stablehlo.constant dense<4> : tensor<i64>
    %2508 = stablehlo.broadcast_in_dim %c_245, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2509 = stablehlo.add %2480, %2508 : tensor<1xi64>
    %2510 = stablehlo.select %2507, %2509, %2480 : tensor<1xi1>, tensor<1xi64>
    %2511 = stablehlo.convert %2510 : (tensor<1xi64>) -> tensor<1xi32>
    %2512 = stablehlo.broadcast_in_dim %2511, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %2513 = stablehlo.broadcast_in_dim %2462, dims = [1] : (tensor<4xf64>) -> tensor<1x4xf64>
    %2514 = "stablehlo.scatter"(%2513, %2512, %2457) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<1x1xi32>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %c_246 = stablehlo.constant dense<0> : tensor<i64>
    %2515 = stablehlo.broadcast_in_dim %c_246, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %2516 = stablehlo.compare  LT, %2505, %2515,  SIGNED : (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi1>
    %c_247 = stablehlo.constant dense<4> : tensor<i64>
    %2517 = stablehlo.broadcast_in_dim %c_247, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %2518 = stablehlo.add %2505, %2517 : tensor<3xi64>
    %2519 = stablehlo.select %2516, %2518, %2505 : tensor<3xi1>, tensor<3xi64>
    %2520 = stablehlo.convert %2519 : (tensor<3xi64>) -> tensor<3xi32>
    %2521 = stablehlo.broadcast_in_dim %2520, dims = [0] : (tensor<3xi32>) -> tensor<3x1xi32>
    %2522 = "stablehlo.scatter"(%2514, %2521, %2456) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<3x1xi32>, tensor<1x3xf64>) -> tensor<1x4xf64>
    %2523 = stablehlo.slice %2389 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2524 = stablehlo.reshape %2523 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2525 = stablehlo.slice %2522 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2526 = stablehlo.reshape %2525 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2527 = stablehlo.multiply %2524, %2526 : tensor<1xf64>
    %2528 = stablehlo.slice %2389 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2529 = stablehlo.reshape %2528 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2530 = stablehlo.slice %2522 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2531 = stablehlo.reshape %2530 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2532 = stablehlo.multiply %2529, %2531 : tensor<1xf64>
    %2533 = stablehlo.subtract %2527, %2532 : tensor<1xf64>
    %2534 = stablehlo.slice %2389 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2535 = stablehlo.reshape %2534 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2536 = stablehlo.slice %2522 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2537 = stablehlo.reshape %2536 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2538 = stablehlo.multiply %2535, %2537 : tensor<1xf64>
    %2539 = stablehlo.subtract %2533, %2538 : tensor<1xf64>
    %2540 = stablehlo.slice %2389 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2541 = stablehlo.reshape %2540 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2542 = stablehlo.slice %2522 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2543 = stablehlo.reshape %2542 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2544 = stablehlo.multiply %2541, %2543 : tensor<1xf64>
    %2545 = stablehlo.subtract %2539, %2544 : tensor<1xf64>
    %2546 = stablehlo.slice %2389 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2547 = stablehlo.reshape %2546 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2548 = stablehlo.slice %2522 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2549 = stablehlo.reshape %2548 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2550 = stablehlo.multiply %2547, %2549 : tensor<1xf64>
    %2551 = stablehlo.slice %2389 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2552 = stablehlo.reshape %2551 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2553 = stablehlo.slice %2522 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2554 = stablehlo.reshape %2553 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2555 = stablehlo.multiply %2552, %2554 : tensor<1xf64>
    %2556 = stablehlo.add %2550, %2555 : tensor<1xf64>
    %2557 = stablehlo.slice %2389 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2558 = stablehlo.reshape %2557 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2559 = stablehlo.slice %2522 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2560 = stablehlo.reshape %2559 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2561 = stablehlo.multiply %2558, %2560 : tensor<1xf64>
    %2562 = stablehlo.add %2556, %2561 : tensor<1xf64>
    %2563 = stablehlo.slice %2389 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2564 = stablehlo.reshape %2563 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2565 = stablehlo.slice %2522 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2566 = stablehlo.reshape %2565 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2567 = stablehlo.multiply %2564, %2566 : tensor<1xf64>
    %2568 = stablehlo.subtract %2562, %2567 : tensor<1xf64>
    %2569 = stablehlo.slice %2389 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2570 = stablehlo.reshape %2569 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2571 = stablehlo.slice %2522 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2572 = stablehlo.reshape %2571 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2573 = stablehlo.multiply %2570, %2572 : tensor<1xf64>
    %2574 = stablehlo.slice %2389 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2575 = stablehlo.reshape %2574 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2576 = stablehlo.slice %2522 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2577 = stablehlo.reshape %2576 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2578 = stablehlo.multiply %2575, %2577 : tensor<1xf64>
    %2579 = stablehlo.subtract %2573, %2578 : tensor<1xf64>
    %2580 = stablehlo.slice %2389 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2581 = stablehlo.reshape %2580 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2582 = stablehlo.slice %2522 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2583 = stablehlo.reshape %2582 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2584 = stablehlo.multiply %2581, %2583 : tensor<1xf64>
    %2585 = stablehlo.add %2579, %2584 : tensor<1xf64>
    %2586 = stablehlo.slice %2389 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2587 = stablehlo.reshape %2586 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2588 = stablehlo.slice %2522 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2589 = stablehlo.reshape %2588 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2590 = stablehlo.multiply %2587, %2589 : tensor<1xf64>
    %2591 = stablehlo.add %2585, %2590 : tensor<1xf64>
    %2592 = stablehlo.slice %2389 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2593 = stablehlo.reshape %2592 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2594 = stablehlo.slice %2522 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2595 = stablehlo.reshape %2594 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2596 = stablehlo.multiply %2593, %2595 : tensor<1xf64>
    %2597 = stablehlo.slice %2389 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2598 = stablehlo.reshape %2597 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2599 = stablehlo.slice %2522 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2600 = stablehlo.reshape %2599 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2601 = stablehlo.multiply %2598, %2600 : tensor<1xf64>
    %2602 = stablehlo.add %2596, %2601 : tensor<1xf64>
    %2603 = stablehlo.slice %2389 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2604 = stablehlo.reshape %2603 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2605 = stablehlo.slice %2522 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2606 = stablehlo.reshape %2605 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2607 = stablehlo.multiply %2604, %2606 : tensor<1xf64>
    %2608 = stablehlo.subtract %2602, %2607 : tensor<1xf64>
    %2609 = stablehlo.slice %2389 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2610 = stablehlo.reshape %2609 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2611 = stablehlo.slice %2522 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2612 = stablehlo.reshape %2611 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2613 = stablehlo.multiply %2610, %2612 : tensor<1xf64>
    %2614 = stablehlo.add %2608, %2613 : tensor<1xf64>
    %2615 = stablehlo.broadcast_in_dim %2545, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2616 = stablehlo.broadcast_in_dim %2568, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2617 = stablehlo.broadcast_in_dim %2591, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2618 = stablehlo.broadcast_in_dim %2614, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2619 = stablehlo.concatenate %2615, %2616, %2617, %2618, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %2620 = stablehlo.slice %2263 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %2621 = stablehlo.reshape %2620 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %2622 = stablehlo.slice %2619 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2623 = stablehlo.reshape %2622 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2624 = stablehlo.slice %2619 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %2625 = stablehlo.dot_general %2624, %2621, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2626 = stablehlo.broadcast_in_dim %2625, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2627 = stablehlo.broadcast_in_dim %2626, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2628 = stablehlo.multiply %2627, %2624 : tensor<1x3xf64>
    %cst_248 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2629 = stablehlo.broadcast_in_dim %cst_248, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %2630 = stablehlo.multiply %2629, %2628 : tensor<1x3xf64>
    %2631 = stablehlo.multiply %2623, %2623 : tensor<1xf64>
    %2632 = stablehlo.dot_general %2624, %2624, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2633 = stablehlo.subtract %2631, %2632 : tensor<1xf64>
    %2634 = stablehlo.broadcast_in_dim %2633, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2635 = stablehlo.broadcast_in_dim %2634, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2636 = stablehlo.multiply %2635, %2621 : tensor<1x3xf64>
    %2637 = stablehlo.add %2630, %2636 : tensor<1x3xf64>
    %cst_249 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2638 = stablehlo.broadcast_in_dim %cst_249, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2639 = stablehlo.multiply %2638, %2623 : tensor<1xf64>
    %2640 = call @cross(%2624, %2621) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %2641 = stablehlo.broadcast_in_dim %2639, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2642 = stablehlo.broadcast_in_dim %2641, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2643 = stablehlo.multiply %2642, %2640 : tensor<1x3xf64>
    %2644 = stablehlo.add %2637, %2643 : tensor<1x3xf64>
    %2645 = stablehlo.subtract %2415, %2644 : tensor<1x3xf64>
    %2646 = stablehlo.broadcast_in_dim %2415, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %2647 = stablehlo.broadcast_in_dim %2440, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %2648 = stablehlo.broadcast_in_dim %2619, dims = [0, 1] : (tensor<1x4xf64>) -> tensor<1x4x1xf64>
    %2649 = stablehlo.broadcast_in_dim %2619, dims = [0, 2] : (tensor<1x4xf64>) -> tensor<1x1x4xf64>
    %2650 = stablehlo.broadcast_in_dim %2648, dims = [0, 1, 2] : (tensor<1x4x1xf64>) -> tensor<1x4x4xf64>
    %2651 = stablehlo.broadcast_in_dim %2649, dims = [0, 1, 2] : (tensor<1x1x4xf64>) -> tensor<1x4x4xf64>
    %2652 = stablehlo.multiply %2650, %2651 : tensor<1x4x4xf64>
    %2653 = stablehlo.slice %2652 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2654 = stablehlo.reshape %2653 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2655 = stablehlo.slice %2652 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2656 = stablehlo.reshape %2655 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2657 = stablehlo.add %2654, %2656 : tensor<1xf64>
    %2658 = stablehlo.slice %2652 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2659 = stablehlo.reshape %2658 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2660 = stablehlo.subtract %2657, %2659 : tensor<1xf64>
    %2661 = stablehlo.slice %2652 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2662 = stablehlo.reshape %2661 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2663 = stablehlo.subtract %2660, %2662 : tensor<1xf64>
    %2664 = stablehlo.slice %2652 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2665 = stablehlo.reshape %2664 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2666 = stablehlo.slice %2652 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2667 = stablehlo.reshape %2666 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2668 = stablehlo.subtract %2665, %2667 : tensor<1xf64>
    %cst_250 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2669 = stablehlo.broadcast_in_dim %cst_250, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2670 = stablehlo.multiply %2669, %2668 : tensor<1xf64>
    %2671 = stablehlo.slice %2652 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2672 = stablehlo.reshape %2671 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2673 = stablehlo.slice %2652 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2674 = stablehlo.reshape %2673 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2675 = stablehlo.add %2672, %2674 : tensor<1xf64>
    %cst_251 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2676 = stablehlo.broadcast_in_dim %cst_251, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2677 = stablehlo.multiply %2676, %2675 : tensor<1xf64>
    %2678 = stablehlo.slice %2652 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2679 = stablehlo.reshape %2678 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2680 = stablehlo.slice %2652 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2681 = stablehlo.reshape %2680 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2682 = stablehlo.add %2679, %2681 : tensor<1xf64>
    %cst_252 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2683 = stablehlo.broadcast_in_dim %cst_252, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2684 = stablehlo.multiply %2683, %2682 : tensor<1xf64>
    %2685 = stablehlo.slice %2652 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2686 = stablehlo.reshape %2685 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2687 = stablehlo.slice %2652 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2688 = stablehlo.reshape %2687 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2689 = stablehlo.subtract %2686, %2688 : tensor<1xf64>
    %2690 = stablehlo.slice %2652 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2691 = stablehlo.reshape %2690 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2692 = stablehlo.add %2689, %2691 : tensor<1xf64>
    %2693 = stablehlo.slice %2652 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2694 = stablehlo.reshape %2693 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2695 = stablehlo.subtract %2692, %2694 : tensor<1xf64>
    %2696 = stablehlo.slice %2652 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2697 = stablehlo.reshape %2696 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2698 = stablehlo.slice %2652 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2699 = stablehlo.reshape %2698 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2700 = stablehlo.subtract %2697, %2699 : tensor<1xf64>
    %cst_253 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2701 = stablehlo.broadcast_in_dim %cst_253, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2702 = stablehlo.multiply %2701, %2700 : tensor<1xf64>
    %2703 = stablehlo.slice %2652 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2704 = stablehlo.reshape %2703 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2705 = stablehlo.slice %2652 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2706 = stablehlo.reshape %2705 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2707 = stablehlo.subtract %2704, %2706 : tensor<1xf64>
    %cst_254 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2708 = stablehlo.broadcast_in_dim %cst_254, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2709 = stablehlo.multiply %2708, %2707 : tensor<1xf64>
    %2710 = stablehlo.slice %2652 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2711 = stablehlo.reshape %2710 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2712 = stablehlo.slice %2652 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2713 = stablehlo.reshape %2712 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2714 = stablehlo.add %2711, %2713 : tensor<1xf64>
    %cst_255 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2715 = stablehlo.broadcast_in_dim %cst_255, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2716 = stablehlo.multiply %2715, %2714 : tensor<1xf64>
    %2717 = stablehlo.slice %2652 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2718 = stablehlo.reshape %2717 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2719 = stablehlo.slice %2652 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2720 = stablehlo.reshape %2719 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2721 = stablehlo.subtract %2718, %2720 : tensor<1xf64>
    %2722 = stablehlo.slice %2652 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2723 = stablehlo.reshape %2722 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2724 = stablehlo.subtract %2721, %2723 : tensor<1xf64>
    %2725 = stablehlo.slice %2652 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %2726 = stablehlo.reshape %2725 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %2727 = stablehlo.add %2724, %2726 : tensor<1xf64>
    %2728 = stablehlo.broadcast_in_dim %2663, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2729 = stablehlo.broadcast_in_dim %2670, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2730 = stablehlo.broadcast_in_dim %2677, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2731 = stablehlo.concatenate %2728, %2729, %2730, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %2732 = stablehlo.broadcast_in_dim %2684, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2733 = stablehlo.broadcast_in_dim %2695, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2734 = stablehlo.broadcast_in_dim %2702, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2735 = stablehlo.concatenate %2732, %2733, %2734, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %2736 = stablehlo.broadcast_in_dim %2709, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2737 = stablehlo.broadcast_in_dim %2716, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2738 = stablehlo.broadcast_in_dim %2727, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2739 = stablehlo.concatenate %2736, %2737, %2738, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %2740 = stablehlo.broadcast_in_dim %2731, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %2741 = stablehlo.broadcast_in_dim %2735, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %2742 = stablehlo.broadcast_in_dim %2739, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %2743 = stablehlo.concatenate %2740, %2741, %2742, dim = 1 : (tensor<1x1x3xf64>, tensor<1x1x3xf64>, tensor<1x1x3xf64>) -> tensor<1x3x3xf64>
    %2744 = call @_take_2(%2645, %c_26) : (tensor<1x3xf64>, tensor<1xi64>) -> tensor<1x3xf64>
    %2745 = call @_take_3(%2619, %c_26) : (tensor<1x4xf64>, tensor<1xi64>) -> tensor<1x4xf64>
    %2746 = call @_take_4(%arg5, %c_27) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %2747 = call @_take_4(%arg6, %c_28) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %2748 = call @_take_6(%arg13, %c_29) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %2749 = call @_take_6(%arg0, %c_30) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %2750 = stablehlo.slice %arg1 [7:8, 0:3] : (tensor<9x3xf64>) -> tensor<1x3xf64>
    %2751 = stablehlo.slice %arg2 [7:8, 0:4] : (tensor<9x4xf64>) -> tensor<1x4xf64>
    %2752 = stablehlo.slice %2745 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2753 = stablehlo.reshape %2752 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2754 = stablehlo.slice %2745 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %2755 = stablehlo.dot_general %2754, %2750, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2756 = stablehlo.broadcast_in_dim %2755, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2757 = stablehlo.broadcast_in_dim %2756, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2758 = stablehlo.multiply %2757, %2754 : tensor<1x3xf64>
    %cst_256 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2759 = stablehlo.broadcast_in_dim %cst_256, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %2760 = stablehlo.multiply %2759, %2758 : tensor<1x3xf64>
    %2761 = stablehlo.multiply %2753, %2753 : tensor<1xf64>
    %2762 = stablehlo.dot_general %2754, %2754, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2763 = stablehlo.subtract %2761, %2762 : tensor<1xf64>
    %2764 = stablehlo.broadcast_in_dim %2763, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2765 = stablehlo.broadcast_in_dim %2764, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2766 = stablehlo.multiply %2765, %2750 : tensor<1x3xf64>
    %2767 = stablehlo.add %2760, %2766 : tensor<1x3xf64>
    %cst_257 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2768 = stablehlo.broadcast_in_dim %cst_257, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2769 = stablehlo.multiply %2768, %2753 : tensor<1xf64>
    %2770 = call @cross(%2754, %2750) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %2771 = stablehlo.broadcast_in_dim %2769, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2772 = stablehlo.broadcast_in_dim %2771, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2773 = stablehlo.multiply %2772, %2770 : tensor<1x3xf64>
    %2774 = stablehlo.add %2767, %2773 : tensor<1x3xf64>
    %2775 = stablehlo.add %2744, %2774 : tensor<1x3xf64>
    %2776 = stablehlo.slice %2745 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2777 = stablehlo.reshape %2776 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2778 = stablehlo.slice %2751 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2779 = stablehlo.reshape %2778 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2780 = stablehlo.multiply %2777, %2779 : tensor<1xf64>
    %2781 = stablehlo.slice %2745 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2782 = stablehlo.reshape %2781 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2783 = stablehlo.slice %2751 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2784 = stablehlo.reshape %2783 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2785 = stablehlo.multiply %2782, %2784 : tensor<1xf64>
    %2786 = stablehlo.subtract %2780, %2785 : tensor<1xf64>
    %2787 = stablehlo.slice %2745 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2788 = stablehlo.reshape %2787 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2789 = stablehlo.slice %2751 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2790 = stablehlo.reshape %2789 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2791 = stablehlo.multiply %2788, %2790 : tensor<1xf64>
    %2792 = stablehlo.subtract %2786, %2791 : tensor<1xf64>
    %2793 = stablehlo.slice %2745 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2794 = stablehlo.reshape %2793 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2795 = stablehlo.slice %2751 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2796 = stablehlo.reshape %2795 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2797 = stablehlo.multiply %2794, %2796 : tensor<1xf64>
    %2798 = stablehlo.subtract %2792, %2797 : tensor<1xf64>
    %2799 = stablehlo.slice %2745 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2800 = stablehlo.reshape %2799 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2801 = stablehlo.slice %2751 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2802 = stablehlo.reshape %2801 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2803 = stablehlo.multiply %2800, %2802 : tensor<1xf64>
    %2804 = stablehlo.slice %2745 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2805 = stablehlo.reshape %2804 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2806 = stablehlo.slice %2751 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2807 = stablehlo.reshape %2806 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2808 = stablehlo.multiply %2805, %2807 : tensor<1xf64>
    %2809 = stablehlo.add %2803, %2808 : tensor<1xf64>
    %2810 = stablehlo.slice %2745 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2811 = stablehlo.reshape %2810 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2812 = stablehlo.slice %2751 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2813 = stablehlo.reshape %2812 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2814 = stablehlo.multiply %2811, %2813 : tensor<1xf64>
    %2815 = stablehlo.add %2809, %2814 : tensor<1xf64>
    %2816 = stablehlo.slice %2745 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2817 = stablehlo.reshape %2816 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2818 = stablehlo.slice %2751 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2819 = stablehlo.reshape %2818 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2820 = stablehlo.multiply %2817, %2819 : tensor<1xf64>
    %2821 = stablehlo.subtract %2815, %2820 : tensor<1xf64>
    %2822 = stablehlo.slice %2745 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2823 = stablehlo.reshape %2822 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2824 = stablehlo.slice %2751 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2825 = stablehlo.reshape %2824 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2826 = stablehlo.multiply %2823, %2825 : tensor<1xf64>
    %2827 = stablehlo.slice %2745 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2828 = stablehlo.reshape %2827 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2829 = stablehlo.slice %2751 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2830 = stablehlo.reshape %2829 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2831 = stablehlo.multiply %2828, %2830 : tensor<1xf64>
    %2832 = stablehlo.subtract %2826, %2831 : tensor<1xf64>
    %2833 = stablehlo.slice %2745 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2834 = stablehlo.reshape %2833 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2835 = stablehlo.slice %2751 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2836 = stablehlo.reshape %2835 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2837 = stablehlo.multiply %2834, %2836 : tensor<1xf64>
    %2838 = stablehlo.add %2832, %2837 : tensor<1xf64>
    %2839 = stablehlo.slice %2745 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2840 = stablehlo.reshape %2839 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2841 = stablehlo.slice %2751 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2842 = stablehlo.reshape %2841 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2843 = stablehlo.multiply %2840, %2842 : tensor<1xf64>
    %2844 = stablehlo.add %2838, %2843 : tensor<1xf64>
    %2845 = stablehlo.slice %2745 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2846 = stablehlo.reshape %2845 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2847 = stablehlo.slice %2751 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2848 = stablehlo.reshape %2847 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2849 = stablehlo.multiply %2846, %2848 : tensor<1xf64>
    %2850 = stablehlo.slice %2745 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2851 = stablehlo.reshape %2850 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2852 = stablehlo.slice %2751 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2853 = stablehlo.reshape %2852 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2854 = stablehlo.multiply %2851, %2853 : tensor<1xf64>
    %2855 = stablehlo.add %2849, %2854 : tensor<1xf64>
    %2856 = stablehlo.slice %2745 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2857 = stablehlo.reshape %2856 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2858 = stablehlo.slice %2751 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2859 = stablehlo.reshape %2858 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2860 = stablehlo.multiply %2857, %2859 : tensor<1xf64>
    %2861 = stablehlo.subtract %2855, %2860 : tensor<1xf64>
    %2862 = stablehlo.slice %2745 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2863 = stablehlo.reshape %2862 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2864 = stablehlo.slice %2751 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2865 = stablehlo.reshape %2864 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2866 = stablehlo.multiply %2863, %2865 : tensor<1xf64>
    %2867 = stablehlo.add %2861, %2866 : tensor<1xf64>
    %2868 = stablehlo.broadcast_in_dim %2798, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2869 = stablehlo.broadcast_in_dim %2821, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2870 = stablehlo.broadcast_in_dim %2844, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2871 = stablehlo.broadcast_in_dim %2867, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2872 = stablehlo.concatenate %2868, %2869, %2870, %2871, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %2873 = stablehlo.slice %2746 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %2874 = stablehlo.reshape %2873 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %2875 = stablehlo.slice %2872 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2876 = stablehlo.reshape %2875 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2877 = stablehlo.slice %2872 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %2878 = stablehlo.dot_general %2877, %2874, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2879 = stablehlo.broadcast_in_dim %2878, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2880 = stablehlo.broadcast_in_dim %2879, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2881 = stablehlo.multiply %2880, %2877 : tensor<1x3xf64>
    %cst_258 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2882 = stablehlo.broadcast_in_dim %cst_258, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %2883 = stablehlo.multiply %2882, %2881 : tensor<1x3xf64>
    %2884 = stablehlo.multiply %2876, %2876 : tensor<1xf64>
    %2885 = stablehlo.dot_general %2877, %2877, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2886 = stablehlo.subtract %2884, %2885 : tensor<1xf64>
    %2887 = stablehlo.broadcast_in_dim %2886, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2888 = stablehlo.broadcast_in_dim %2887, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2889 = stablehlo.multiply %2888, %2874 : tensor<1x3xf64>
    %2890 = stablehlo.add %2883, %2889 : tensor<1x3xf64>
    %cst_259 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2891 = stablehlo.broadcast_in_dim %cst_259, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2892 = stablehlo.multiply %2891, %2876 : tensor<1xf64>
    %2893 = call @cross(%2877, %2874) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %2894 = stablehlo.broadcast_in_dim %2892, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2895 = stablehlo.broadcast_in_dim %2894, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2896 = stablehlo.multiply %2895, %2893 : tensor<1x3xf64>
    %2897 = stablehlo.add %2890, %2896 : tensor<1x3xf64>
    %2898 = stablehlo.add %2897, %2775 : tensor<1x3xf64>
    %2899 = stablehlo.slice %2747 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %2900 = stablehlo.reshape %2899 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %2901 = stablehlo.slice %2872 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %2902 = stablehlo.reshape %2901 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2903 = stablehlo.slice %2872 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %2904 = stablehlo.dot_general %2903, %2900, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2905 = stablehlo.broadcast_in_dim %2904, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2906 = stablehlo.broadcast_in_dim %2905, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2907 = stablehlo.multiply %2906, %2903 : tensor<1x3xf64>
    %cst_260 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2908 = stablehlo.broadcast_in_dim %cst_260, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %2909 = stablehlo.multiply %2908, %2907 : tensor<1x3xf64>
    %2910 = stablehlo.multiply %2902, %2902 : tensor<1xf64>
    %2911 = stablehlo.dot_general %2903, %2903, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %2912 = stablehlo.subtract %2910, %2911 : tensor<1xf64>
    %2913 = stablehlo.broadcast_in_dim %2912, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2914 = stablehlo.broadcast_in_dim %2913, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2915 = stablehlo.multiply %2914, %2900 : tensor<1x3xf64>
    %2916 = stablehlo.add %2909, %2915 : tensor<1x3xf64>
    %cst_261 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %2917 = stablehlo.broadcast_in_dim %cst_261, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2918 = stablehlo.multiply %2917, %2902 : tensor<1xf64>
    %2919 = call @cross(%2903, %2900) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %2920 = stablehlo.broadcast_in_dim %2918, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2921 = stablehlo.broadcast_in_dim %2920, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2922 = stablehlo.multiply %2921, %2919 : tensor<1x3xf64>
    %2923 = stablehlo.add %2916, %2922 : tensor<1x3xf64>
    %2924 = stablehlo.slice %2748 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %2925 = stablehlo.reshape %2924 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2926 = stablehlo.slice %2749 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %2927 = stablehlo.reshape %2926 : (tensor<1x1xf64>) -> tensor<1xf64>
    %2928 = stablehlo.subtract %2925, %2927 : tensor<1xf64>
    %2929 = stablehlo.slice %2747 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %2930 = stablehlo.reshape %2929 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %cst_262 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %2931 = stablehlo.broadcast_in_dim %cst_262, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2932 = stablehlo.multiply %2928, %2931 : tensor<1xf64>
    %2933 = stablehlo.sine %2932 : tensor<1xf64>
    %cst_263 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %2934 = stablehlo.broadcast_in_dim %cst_263, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %2935 = stablehlo.multiply %2928, %2934 : tensor<1xf64>
    %2936 = stablehlo.cosine %2935 : tensor<1xf64>
    %2937 = stablehlo.broadcast_in_dim %2933, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %2938 = stablehlo.broadcast_in_dim %2937, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %2939 = stablehlo.multiply %2930, %2938 : tensor<1x3xf64>
    %2940 = stablehlo.broadcast_in_dim %2936, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %c_264 = stablehlo.constant dense<0> : tensor<i64>
    %2941 = stablehlo.reshape %c_264 : (tensor<i64>) -> tensor<1xi64>
    %2942 = stablehlo.slice %2941 [0:1] : (tensor<1xi64>) -> tensor<1xi64>
    %2943 = stablehlo.reshape %2942 : (tensor<1xi64>) -> tensor<i64>
    %2944 = stablehlo.broadcast_in_dim %2943, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %cst_265 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %2945 = stablehlo.broadcast_in_dim %cst_265, dims = [] : (tensor<f64>) -> tensor<4xf64>
    %c_266 = stablehlo.constant dense<0> : tensor<i64>
    %2946 = stablehlo.broadcast_in_dim %c_266, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2947 = stablehlo.compare  LT, %2944, %2946,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_267 = stablehlo.constant dense<3> : tensor<i64>
    %2948 = stablehlo.broadcast_in_dim %c_267, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2949 = stablehlo.add %2944, %2948 : tensor<1xi64>
    %2950 = call @_where_7(%2947, %2949, %2944) : (tensor<1xi1>, tensor<1xi64>, tensor<1xi64>) -> tensor<1xi64>
    %c_268 = stablehlo.constant dense<0> : tensor<i64>
    %c_269 = stablehlo.constant dense<3> : tensor<i64>
    %2951 = call @clip(%2950, %c_268, %c_269) : (tensor<1xi64>, tensor<i64>, tensor<i64>) -> tensor<1xi64>
    %2952 = call @argsort(%2951) : (tensor<1xi64>) -> tensor<1xi64>
    %2953 = stablehlo.iota dim = 0 : tensor<1xi64>
    %c_270 = stablehlo.constant dense<0> : tensor<i64>
    %2954 = stablehlo.broadcast_in_dim %c_270, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2955 = stablehlo.compare  LT, %2952, %2954,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_271 = stablehlo.constant dense<1> : tensor<i64>
    %2956 = stablehlo.broadcast_in_dim %c_271, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2957 = stablehlo.add %2952, %2956 : tensor<1xi64>
    %2958 = stablehlo.select %2955, %2957, %2952 : tensor<1xi1>, tensor<1xi64>
    %2959 = stablehlo.convert %2958 : (tensor<1xi64>) -> tensor<1xi32>
    %2960 = stablehlo.broadcast_in_dim %2959, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %2961 = stablehlo.convert %2951 : tensor<1xi64>
    %2962 = "stablehlo.scatter"(%2961, %2960, %2953) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<1xi64>, tensor<1x1xi32>, tensor<1xi64>) -> tensor<1xi64>
    %2963 = stablehlo.convert %2962 : tensor<1xi64>
    %c_272 = stablehlo.constant dense<true> : tensor<i1>
    %2964 = stablehlo.broadcast_in_dim %c_272, dims = [] : (tensor<i1>) -> tensor<4xi1>
    %c_273 = stablehlo.constant dense<0> : tensor<i64>
    %2965 = stablehlo.broadcast_in_dim %c_273, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2966 = stablehlo.compare  LT, %2963, %2965,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_274 = stablehlo.constant dense<4> : tensor<i64>
    %2967 = stablehlo.broadcast_in_dim %c_274, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2968 = stablehlo.add %2963, %2967 : tensor<1xi64>
    %2969 = stablehlo.select %2966, %2968, %2963 : tensor<1xi1>, tensor<1xi64>
    %2970 = stablehlo.convert %2969 : (tensor<1xi64>) -> tensor<1xi32>
    %2971 = stablehlo.broadcast_in_dim %2970, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %c_275 = stablehlo.constant dense<false> : tensor<i1>
    %2972 = stablehlo.broadcast_in_dim %c_275, dims = [] : (tensor<i1>) -> tensor<1xi1>
    %2973 = "stablehlo.scatter"(%2964, %2971, %2972) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i1>, %arg114: tensor<i1>):
      stablehlo.return %arg114 : tensor<i1>
    }) : (tensor<4xi1>, tensor<1x1xi32>, tensor<1xi1>) -> tensor<4xi1>
    %2974 = call @_cumulative_reduction(%2973) : (tensor<4xi1>) -> tensor<4xi64>
    %c_276 = stablehlo.constant dense<0> : tensor<i64>
    %2975 = stablehlo.broadcast_in_dim %c_276, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %c_277 = stablehlo.constant dense<0> : tensor<i64>
    %2976 = call @clip_8(%2974, %c_277) : (tensor<4xi64>, tensor<i64>) -> tensor<4xi64>
    %c_278 = stablehlo.constant dense<1> : tensor<i64>
    %c_279 = stablehlo.constant dense<0> : tensor<i64>
    %2977 = stablehlo.broadcast_in_dim %c_279, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %2978 = stablehlo.compare  LT, %2976, %2977,  SIGNED : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1>
    %c_280 = stablehlo.constant dense<3> : tensor<i64>
    %2979 = stablehlo.broadcast_in_dim %c_280, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %2980 = stablehlo.add %2976, %2979 : tensor<4xi64>
    %2981 = stablehlo.select %2978, %2980, %2976 : tensor<4xi1>, tensor<4xi64>
    %2982 = stablehlo.convert %2981 : (tensor<4xi64>) -> tensor<4xi32>
    %2983 = stablehlo.broadcast_in_dim %2982, dims = [0] : (tensor<4xi32>) -> tensor<4x1xi32>
    %2984 = stablehlo.broadcast_in_dim %c_278, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %2985 = "stablehlo.scatter"(%2975, %2983, %2984) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<3xi64>, tensor<4x1xi32>, tensor<4xi64>) -> tensor<3xi64>
    %2986 = call @_cumulative_reduction_9(%2985) : (tensor<3xi64>) -> tensor<3xi64>
    %c_281 = stablehlo.constant dense<1> : tensor<i64>
    %2987 = call @floor_divide(%2986, %c_281) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_282 = stablehlo.constant dense<4> : tensor<i64>
    %2988 = call @remainder_12(%2987, %c_282) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_283 = stablehlo.constant dense<0> : tensor<i64>
    %2989 = stablehlo.broadcast_in_dim %c_283, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2990 = stablehlo.compare  LT, %2963, %2989,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_284 = stablehlo.constant dense<4> : tensor<i64>
    %2991 = stablehlo.broadcast_in_dim %c_284, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %2992 = stablehlo.add %2963, %2991 : tensor<1xi64>
    %2993 = stablehlo.select %2990, %2992, %2963 : tensor<1xi1>, tensor<1xi64>
    %2994 = stablehlo.convert %2993 : (tensor<1xi64>) -> tensor<1xi32>
    %2995 = stablehlo.broadcast_in_dim %2994, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %2996 = stablehlo.broadcast_in_dim %2945, dims = [1] : (tensor<4xf64>) -> tensor<1x4xf64>
    %2997 = "stablehlo.scatter"(%2996, %2995, %2940) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<1x1xi32>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %c_285 = stablehlo.constant dense<0> : tensor<i64>
    %2998 = stablehlo.broadcast_in_dim %c_285, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %2999 = stablehlo.compare  LT, %2988, %2998,  SIGNED : (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi1>
    %c_286 = stablehlo.constant dense<4> : tensor<i64>
    %3000 = stablehlo.broadcast_in_dim %c_286, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %3001 = stablehlo.add %2988, %3000 : tensor<3xi64>
    %3002 = stablehlo.select %2999, %3001, %2988 : tensor<3xi1>, tensor<3xi64>
    %3003 = stablehlo.convert %3002 : (tensor<3xi64>) -> tensor<3xi32>
    %3004 = stablehlo.broadcast_in_dim %3003, dims = [0] : (tensor<3xi32>) -> tensor<3x1xi32>
    %3005 = "stablehlo.scatter"(%2997, %3004, %2939) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<3x1xi32>, tensor<1x3xf64>) -> tensor<1x4xf64>
    %3006 = stablehlo.slice %2872 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3007 = stablehlo.reshape %3006 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3008 = stablehlo.slice %3005 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3009 = stablehlo.reshape %3008 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3010 = stablehlo.multiply %3007, %3009 : tensor<1xf64>
    %3011 = stablehlo.slice %2872 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3012 = stablehlo.reshape %3011 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3013 = stablehlo.slice %3005 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3014 = stablehlo.reshape %3013 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3015 = stablehlo.multiply %3012, %3014 : tensor<1xf64>
    %3016 = stablehlo.subtract %3010, %3015 : tensor<1xf64>
    %3017 = stablehlo.slice %2872 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3018 = stablehlo.reshape %3017 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3019 = stablehlo.slice %3005 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3020 = stablehlo.reshape %3019 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3021 = stablehlo.multiply %3018, %3020 : tensor<1xf64>
    %3022 = stablehlo.subtract %3016, %3021 : tensor<1xf64>
    %3023 = stablehlo.slice %2872 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3024 = stablehlo.reshape %3023 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3025 = stablehlo.slice %3005 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3026 = stablehlo.reshape %3025 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3027 = stablehlo.multiply %3024, %3026 : tensor<1xf64>
    %3028 = stablehlo.subtract %3022, %3027 : tensor<1xf64>
    %3029 = stablehlo.slice %2872 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3030 = stablehlo.reshape %3029 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3031 = stablehlo.slice %3005 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3032 = stablehlo.reshape %3031 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3033 = stablehlo.multiply %3030, %3032 : tensor<1xf64>
    %3034 = stablehlo.slice %2872 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3035 = stablehlo.reshape %3034 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3036 = stablehlo.slice %3005 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3037 = stablehlo.reshape %3036 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3038 = stablehlo.multiply %3035, %3037 : tensor<1xf64>
    %3039 = stablehlo.add %3033, %3038 : tensor<1xf64>
    %3040 = stablehlo.slice %2872 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3041 = stablehlo.reshape %3040 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3042 = stablehlo.slice %3005 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3043 = stablehlo.reshape %3042 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3044 = stablehlo.multiply %3041, %3043 : tensor<1xf64>
    %3045 = stablehlo.add %3039, %3044 : tensor<1xf64>
    %3046 = stablehlo.slice %2872 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3047 = stablehlo.reshape %3046 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3048 = stablehlo.slice %3005 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3049 = stablehlo.reshape %3048 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3050 = stablehlo.multiply %3047, %3049 : tensor<1xf64>
    %3051 = stablehlo.subtract %3045, %3050 : tensor<1xf64>
    %3052 = stablehlo.slice %2872 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3053 = stablehlo.reshape %3052 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3054 = stablehlo.slice %3005 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3055 = stablehlo.reshape %3054 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3056 = stablehlo.multiply %3053, %3055 : tensor<1xf64>
    %3057 = stablehlo.slice %2872 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3058 = stablehlo.reshape %3057 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3059 = stablehlo.slice %3005 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3060 = stablehlo.reshape %3059 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3061 = stablehlo.multiply %3058, %3060 : tensor<1xf64>
    %3062 = stablehlo.subtract %3056, %3061 : tensor<1xf64>
    %3063 = stablehlo.slice %2872 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3064 = stablehlo.reshape %3063 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3065 = stablehlo.slice %3005 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3066 = stablehlo.reshape %3065 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3067 = stablehlo.multiply %3064, %3066 : tensor<1xf64>
    %3068 = stablehlo.add %3062, %3067 : tensor<1xf64>
    %3069 = stablehlo.slice %2872 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3070 = stablehlo.reshape %3069 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3071 = stablehlo.slice %3005 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3072 = stablehlo.reshape %3071 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3073 = stablehlo.multiply %3070, %3072 : tensor<1xf64>
    %3074 = stablehlo.add %3068, %3073 : tensor<1xf64>
    %3075 = stablehlo.slice %2872 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3076 = stablehlo.reshape %3075 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3077 = stablehlo.slice %3005 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3078 = stablehlo.reshape %3077 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3079 = stablehlo.multiply %3076, %3078 : tensor<1xf64>
    %3080 = stablehlo.slice %2872 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3081 = stablehlo.reshape %3080 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3082 = stablehlo.slice %3005 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3083 = stablehlo.reshape %3082 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3084 = stablehlo.multiply %3081, %3083 : tensor<1xf64>
    %3085 = stablehlo.add %3079, %3084 : tensor<1xf64>
    %3086 = stablehlo.slice %2872 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3087 = stablehlo.reshape %3086 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3088 = stablehlo.slice %3005 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3089 = stablehlo.reshape %3088 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3090 = stablehlo.multiply %3087, %3089 : tensor<1xf64>
    %3091 = stablehlo.subtract %3085, %3090 : tensor<1xf64>
    %3092 = stablehlo.slice %2872 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3093 = stablehlo.reshape %3092 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3094 = stablehlo.slice %3005 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3095 = stablehlo.reshape %3094 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3096 = stablehlo.multiply %3093, %3095 : tensor<1xf64>
    %3097 = stablehlo.add %3091, %3096 : tensor<1xf64>
    %3098 = stablehlo.broadcast_in_dim %3028, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3099 = stablehlo.broadcast_in_dim %3051, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3100 = stablehlo.broadcast_in_dim %3074, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3101 = stablehlo.broadcast_in_dim %3097, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3102 = stablehlo.concatenate %3098, %3099, %3100, %3101, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %3103 = stablehlo.slice %2746 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %3104 = stablehlo.reshape %3103 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3105 = stablehlo.slice %3102 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3106 = stablehlo.reshape %3105 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3107 = stablehlo.slice %3102 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %3108 = stablehlo.dot_general %3107, %3104, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %3109 = stablehlo.broadcast_in_dim %3108, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3110 = stablehlo.broadcast_in_dim %3109, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3111 = stablehlo.multiply %3110, %3107 : tensor<1x3xf64>
    %cst_287 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3112 = stablehlo.broadcast_in_dim %cst_287, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %3113 = stablehlo.multiply %3112, %3111 : tensor<1x3xf64>
    %3114 = stablehlo.multiply %3106, %3106 : tensor<1xf64>
    %3115 = stablehlo.dot_general %3107, %3107, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %3116 = stablehlo.subtract %3114, %3115 : tensor<1xf64>
    %3117 = stablehlo.broadcast_in_dim %3116, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3118 = stablehlo.broadcast_in_dim %3117, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3119 = stablehlo.multiply %3118, %3104 : tensor<1x3xf64>
    %3120 = stablehlo.add %3113, %3119 : tensor<1x3xf64>
    %cst_288 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3121 = stablehlo.broadcast_in_dim %cst_288, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3122 = stablehlo.multiply %3121, %3106 : tensor<1xf64>
    %3123 = call @cross(%3107, %3104) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %3124 = stablehlo.broadcast_in_dim %3122, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3125 = stablehlo.broadcast_in_dim %3124, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3126 = stablehlo.multiply %3125, %3123 : tensor<1x3xf64>
    %3127 = stablehlo.add %3120, %3126 : tensor<1x3xf64>
    %3128 = stablehlo.subtract %2898, %3127 : tensor<1x3xf64>
    %3129 = stablehlo.broadcast_in_dim %2898, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %3130 = stablehlo.broadcast_in_dim %2923, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %3131 = stablehlo.broadcast_in_dim %3102, dims = [0, 1] : (tensor<1x4xf64>) -> tensor<1x4x1xf64>
    %3132 = stablehlo.broadcast_in_dim %3102, dims = [0, 2] : (tensor<1x4xf64>) -> tensor<1x1x4xf64>
    %3133 = stablehlo.broadcast_in_dim %3131, dims = [0, 1, 2] : (tensor<1x4x1xf64>) -> tensor<1x4x4xf64>
    %3134 = stablehlo.broadcast_in_dim %3132, dims = [0, 1, 2] : (tensor<1x1x4xf64>) -> tensor<1x4x4xf64>
    %3135 = stablehlo.multiply %3133, %3134 : tensor<1x4x4xf64>
    %3136 = stablehlo.slice %3135 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3137 = stablehlo.reshape %3136 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3138 = stablehlo.slice %3135 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3139 = stablehlo.reshape %3138 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3140 = stablehlo.add %3137, %3139 : tensor<1xf64>
    %3141 = stablehlo.slice %3135 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3142 = stablehlo.reshape %3141 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3143 = stablehlo.subtract %3140, %3142 : tensor<1xf64>
    %3144 = stablehlo.slice %3135 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3145 = stablehlo.reshape %3144 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3146 = stablehlo.subtract %3143, %3145 : tensor<1xf64>
    %3147 = stablehlo.slice %3135 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3148 = stablehlo.reshape %3147 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3149 = stablehlo.slice %3135 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3150 = stablehlo.reshape %3149 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3151 = stablehlo.subtract %3148, %3150 : tensor<1xf64>
    %cst_289 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3152 = stablehlo.broadcast_in_dim %cst_289, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3153 = stablehlo.multiply %3152, %3151 : tensor<1xf64>
    %3154 = stablehlo.slice %3135 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3155 = stablehlo.reshape %3154 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3156 = stablehlo.slice %3135 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3157 = stablehlo.reshape %3156 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3158 = stablehlo.add %3155, %3157 : tensor<1xf64>
    %cst_290 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3159 = stablehlo.broadcast_in_dim %cst_290, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3160 = stablehlo.multiply %3159, %3158 : tensor<1xf64>
    %3161 = stablehlo.slice %3135 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3162 = stablehlo.reshape %3161 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3163 = stablehlo.slice %3135 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3164 = stablehlo.reshape %3163 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3165 = stablehlo.add %3162, %3164 : tensor<1xf64>
    %cst_291 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3166 = stablehlo.broadcast_in_dim %cst_291, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3167 = stablehlo.multiply %3166, %3165 : tensor<1xf64>
    %3168 = stablehlo.slice %3135 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3169 = stablehlo.reshape %3168 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3170 = stablehlo.slice %3135 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3171 = stablehlo.reshape %3170 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3172 = stablehlo.subtract %3169, %3171 : tensor<1xf64>
    %3173 = stablehlo.slice %3135 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3174 = stablehlo.reshape %3173 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3175 = stablehlo.add %3172, %3174 : tensor<1xf64>
    %3176 = stablehlo.slice %3135 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3177 = stablehlo.reshape %3176 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3178 = stablehlo.subtract %3175, %3177 : tensor<1xf64>
    %3179 = stablehlo.slice %3135 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3180 = stablehlo.reshape %3179 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3181 = stablehlo.slice %3135 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3182 = stablehlo.reshape %3181 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3183 = stablehlo.subtract %3180, %3182 : tensor<1xf64>
    %cst_292 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3184 = stablehlo.broadcast_in_dim %cst_292, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3185 = stablehlo.multiply %3184, %3183 : tensor<1xf64>
    %3186 = stablehlo.slice %3135 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3187 = stablehlo.reshape %3186 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3188 = stablehlo.slice %3135 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3189 = stablehlo.reshape %3188 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3190 = stablehlo.subtract %3187, %3189 : tensor<1xf64>
    %cst_293 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3191 = stablehlo.broadcast_in_dim %cst_293, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3192 = stablehlo.multiply %3191, %3190 : tensor<1xf64>
    %3193 = stablehlo.slice %3135 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3194 = stablehlo.reshape %3193 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3195 = stablehlo.slice %3135 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3196 = stablehlo.reshape %3195 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3197 = stablehlo.add %3194, %3196 : tensor<1xf64>
    %cst_294 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3198 = stablehlo.broadcast_in_dim %cst_294, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3199 = stablehlo.multiply %3198, %3197 : tensor<1xf64>
    %3200 = stablehlo.slice %3135 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3201 = stablehlo.reshape %3200 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3202 = stablehlo.slice %3135 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3203 = stablehlo.reshape %3202 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3204 = stablehlo.subtract %3201, %3203 : tensor<1xf64>
    %3205 = stablehlo.slice %3135 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3206 = stablehlo.reshape %3205 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3207 = stablehlo.subtract %3204, %3206 : tensor<1xf64>
    %3208 = stablehlo.slice %3135 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3209 = stablehlo.reshape %3208 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3210 = stablehlo.add %3207, %3209 : tensor<1xf64>
    %3211 = stablehlo.broadcast_in_dim %3146, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3212 = stablehlo.broadcast_in_dim %3153, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3213 = stablehlo.broadcast_in_dim %3160, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3214 = stablehlo.concatenate %3211, %3212, %3213, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %3215 = stablehlo.broadcast_in_dim %3167, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3216 = stablehlo.broadcast_in_dim %3178, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3217 = stablehlo.broadcast_in_dim %3185, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3218 = stablehlo.concatenate %3215, %3216, %3217, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %3219 = stablehlo.broadcast_in_dim %3192, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3220 = stablehlo.broadcast_in_dim %3199, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3221 = stablehlo.broadcast_in_dim %3210, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3222 = stablehlo.concatenate %3219, %3220, %3221, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %3223 = stablehlo.broadcast_in_dim %3214, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %3224 = stablehlo.broadcast_in_dim %3218, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %3225 = stablehlo.broadcast_in_dim %3222, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %3226 = stablehlo.concatenate %3223, %3224, %3225, dim = 1 : (tensor<1x1x3xf64>, tensor<1x1x3xf64>, tensor<1x1x3xf64>) -> tensor<1x3x3xf64>
    %3227 = call @_take_2(%3128, %c_31) : (tensor<1x3xf64>, tensor<1xi64>) -> tensor<1x3xf64>
    %3228 = call @_take_3(%3102, %c_31) : (tensor<1x4xf64>, tensor<1xi64>) -> tensor<1x4xf64>
    %3229 = call @_take_4(%arg5, %c_32) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %3230 = call @_take_4(%arg6, %c_33) : (tensor<7x3xf64>, tensor<1x1xi64>) -> tensor<1x1x3xf64>
    %3231 = call @_take_6(%arg13, %c_34) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %3232 = call @_take_6(%arg0, %c_35) : (tensor<7xf64>, tensor<1x1xi64>) -> tensor<1x1xf64>
    %3233 = stablehlo.slice %arg1 [8:9, 0:3] : (tensor<9x3xf64>) -> tensor<1x3xf64>
    %3234 = stablehlo.slice %arg2 [8:9, 0:4] : (tensor<9x4xf64>) -> tensor<1x4xf64>
    %3235 = stablehlo.slice %3228 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3236 = stablehlo.reshape %3235 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3237 = stablehlo.slice %3228 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %3238 = stablehlo.dot_general %3237, %3233, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %3239 = stablehlo.broadcast_in_dim %3238, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3240 = stablehlo.broadcast_in_dim %3239, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3241 = stablehlo.multiply %3240, %3237 : tensor<1x3xf64>
    %cst_295 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3242 = stablehlo.broadcast_in_dim %cst_295, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %3243 = stablehlo.multiply %3242, %3241 : tensor<1x3xf64>
    %3244 = stablehlo.multiply %3236, %3236 : tensor<1xf64>
    %3245 = stablehlo.dot_general %3237, %3237, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %3246 = stablehlo.subtract %3244, %3245 : tensor<1xf64>
    %3247 = stablehlo.broadcast_in_dim %3246, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3248 = stablehlo.broadcast_in_dim %3247, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3249 = stablehlo.multiply %3248, %3233 : tensor<1x3xf64>
    %3250 = stablehlo.add %3243, %3249 : tensor<1x3xf64>
    %cst_296 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3251 = stablehlo.broadcast_in_dim %cst_296, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3252 = stablehlo.multiply %3251, %3236 : tensor<1xf64>
    %3253 = call @cross(%3237, %3233) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %3254 = stablehlo.broadcast_in_dim %3252, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3255 = stablehlo.broadcast_in_dim %3254, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3256 = stablehlo.multiply %3255, %3253 : tensor<1x3xf64>
    %3257 = stablehlo.add %3250, %3256 : tensor<1x3xf64>
    %3258 = stablehlo.add %3227, %3257 : tensor<1x3xf64>
    %3259 = stablehlo.slice %3228 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3260 = stablehlo.reshape %3259 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3261 = stablehlo.slice %3234 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3262 = stablehlo.reshape %3261 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3263 = stablehlo.multiply %3260, %3262 : tensor<1xf64>
    %3264 = stablehlo.slice %3228 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3265 = stablehlo.reshape %3264 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3266 = stablehlo.slice %3234 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3267 = stablehlo.reshape %3266 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3268 = stablehlo.multiply %3265, %3267 : tensor<1xf64>
    %3269 = stablehlo.subtract %3263, %3268 : tensor<1xf64>
    %3270 = stablehlo.slice %3228 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3271 = stablehlo.reshape %3270 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3272 = stablehlo.slice %3234 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3273 = stablehlo.reshape %3272 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3274 = stablehlo.multiply %3271, %3273 : tensor<1xf64>
    %3275 = stablehlo.subtract %3269, %3274 : tensor<1xf64>
    %3276 = stablehlo.slice %3228 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3277 = stablehlo.reshape %3276 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3278 = stablehlo.slice %3234 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3279 = stablehlo.reshape %3278 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3280 = stablehlo.multiply %3277, %3279 : tensor<1xf64>
    %3281 = stablehlo.subtract %3275, %3280 : tensor<1xf64>
    %3282 = stablehlo.slice %3228 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3283 = stablehlo.reshape %3282 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3284 = stablehlo.slice %3234 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3285 = stablehlo.reshape %3284 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3286 = stablehlo.multiply %3283, %3285 : tensor<1xf64>
    %3287 = stablehlo.slice %3228 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3288 = stablehlo.reshape %3287 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3289 = stablehlo.slice %3234 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3290 = stablehlo.reshape %3289 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3291 = stablehlo.multiply %3288, %3290 : tensor<1xf64>
    %3292 = stablehlo.add %3286, %3291 : tensor<1xf64>
    %3293 = stablehlo.slice %3228 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3294 = stablehlo.reshape %3293 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3295 = stablehlo.slice %3234 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3296 = stablehlo.reshape %3295 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3297 = stablehlo.multiply %3294, %3296 : tensor<1xf64>
    %3298 = stablehlo.add %3292, %3297 : tensor<1xf64>
    %3299 = stablehlo.slice %3228 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3300 = stablehlo.reshape %3299 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3301 = stablehlo.slice %3234 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3302 = stablehlo.reshape %3301 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3303 = stablehlo.multiply %3300, %3302 : tensor<1xf64>
    %3304 = stablehlo.subtract %3298, %3303 : tensor<1xf64>
    %3305 = stablehlo.slice %3228 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3306 = stablehlo.reshape %3305 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3307 = stablehlo.slice %3234 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3308 = stablehlo.reshape %3307 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3309 = stablehlo.multiply %3306, %3308 : tensor<1xf64>
    %3310 = stablehlo.slice %3228 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3311 = stablehlo.reshape %3310 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3312 = stablehlo.slice %3234 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3313 = stablehlo.reshape %3312 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3314 = stablehlo.multiply %3311, %3313 : tensor<1xf64>
    %3315 = stablehlo.subtract %3309, %3314 : tensor<1xf64>
    %3316 = stablehlo.slice %3228 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3317 = stablehlo.reshape %3316 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3318 = stablehlo.slice %3234 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3319 = stablehlo.reshape %3318 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3320 = stablehlo.multiply %3317, %3319 : tensor<1xf64>
    %3321 = stablehlo.add %3315, %3320 : tensor<1xf64>
    %3322 = stablehlo.slice %3228 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3323 = stablehlo.reshape %3322 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3324 = stablehlo.slice %3234 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3325 = stablehlo.reshape %3324 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3326 = stablehlo.multiply %3323, %3325 : tensor<1xf64>
    %3327 = stablehlo.add %3321, %3326 : tensor<1xf64>
    %3328 = stablehlo.slice %3228 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3329 = stablehlo.reshape %3328 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3330 = stablehlo.slice %3234 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3331 = stablehlo.reshape %3330 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3332 = stablehlo.multiply %3329, %3331 : tensor<1xf64>
    %3333 = stablehlo.slice %3228 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3334 = stablehlo.reshape %3333 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3335 = stablehlo.slice %3234 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3336 = stablehlo.reshape %3335 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3337 = stablehlo.multiply %3334, %3336 : tensor<1xf64>
    %3338 = stablehlo.add %3332, %3337 : tensor<1xf64>
    %3339 = stablehlo.slice %3228 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3340 = stablehlo.reshape %3339 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3341 = stablehlo.slice %3234 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3342 = stablehlo.reshape %3341 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3343 = stablehlo.multiply %3340, %3342 : tensor<1xf64>
    %3344 = stablehlo.subtract %3338, %3343 : tensor<1xf64>
    %3345 = stablehlo.slice %3228 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3346 = stablehlo.reshape %3345 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3347 = stablehlo.slice %3234 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3348 = stablehlo.reshape %3347 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3349 = stablehlo.multiply %3346, %3348 : tensor<1xf64>
    %3350 = stablehlo.add %3344, %3349 : tensor<1xf64>
    %3351 = stablehlo.broadcast_in_dim %3281, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3352 = stablehlo.broadcast_in_dim %3304, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3353 = stablehlo.broadcast_in_dim %3327, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3354 = stablehlo.broadcast_in_dim %3350, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3355 = stablehlo.concatenate %3351, %3352, %3353, %3354, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %3356 = stablehlo.slice %3229 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %3357 = stablehlo.reshape %3356 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3358 = stablehlo.slice %3355 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3359 = stablehlo.reshape %3358 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3360 = stablehlo.slice %3355 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %3361 = stablehlo.dot_general %3360, %3357, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %3362 = stablehlo.broadcast_in_dim %3361, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3363 = stablehlo.broadcast_in_dim %3362, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3364 = stablehlo.multiply %3363, %3360 : tensor<1x3xf64>
    %cst_297 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3365 = stablehlo.broadcast_in_dim %cst_297, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %3366 = stablehlo.multiply %3365, %3364 : tensor<1x3xf64>
    %3367 = stablehlo.multiply %3359, %3359 : tensor<1xf64>
    %3368 = stablehlo.dot_general %3360, %3360, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %3369 = stablehlo.subtract %3367, %3368 : tensor<1xf64>
    %3370 = stablehlo.broadcast_in_dim %3369, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3371 = stablehlo.broadcast_in_dim %3370, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3372 = stablehlo.multiply %3371, %3357 : tensor<1x3xf64>
    %3373 = stablehlo.add %3366, %3372 : tensor<1x3xf64>
    %cst_298 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3374 = stablehlo.broadcast_in_dim %cst_298, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3375 = stablehlo.multiply %3374, %3359 : tensor<1xf64>
    %3376 = call @cross(%3360, %3357) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %3377 = stablehlo.broadcast_in_dim %3375, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3378 = stablehlo.broadcast_in_dim %3377, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3379 = stablehlo.multiply %3378, %3376 : tensor<1x3xf64>
    %3380 = stablehlo.add %3373, %3379 : tensor<1x3xf64>
    %3381 = stablehlo.add %3380, %3258 : tensor<1x3xf64>
    %3382 = stablehlo.slice %3230 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %3383 = stablehlo.reshape %3382 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3384 = stablehlo.slice %3355 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3385 = stablehlo.reshape %3384 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3386 = stablehlo.slice %3355 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %3387 = stablehlo.dot_general %3386, %3383, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %3388 = stablehlo.broadcast_in_dim %3387, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3389 = stablehlo.broadcast_in_dim %3388, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3390 = stablehlo.multiply %3389, %3386 : tensor<1x3xf64>
    %cst_299 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3391 = stablehlo.broadcast_in_dim %cst_299, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %3392 = stablehlo.multiply %3391, %3390 : tensor<1x3xf64>
    %3393 = stablehlo.multiply %3385, %3385 : tensor<1xf64>
    %3394 = stablehlo.dot_general %3386, %3386, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %3395 = stablehlo.subtract %3393, %3394 : tensor<1xf64>
    %3396 = stablehlo.broadcast_in_dim %3395, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3397 = stablehlo.broadcast_in_dim %3396, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3398 = stablehlo.multiply %3397, %3383 : tensor<1x3xf64>
    %3399 = stablehlo.add %3392, %3398 : tensor<1x3xf64>
    %cst_300 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3400 = stablehlo.broadcast_in_dim %cst_300, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3401 = stablehlo.multiply %3400, %3385 : tensor<1xf64>
    %3402 = call @cross(%3386, %3383) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %3403 = stablehlo.broadcast_in_dim %3401, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3404 = stablehlo.broadcast_in_dim %3403, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3405 = stablehlo.multiply %3404, %3402 : tensor<1x3xf64>
    %3406 = stablehlo.add %3399, %3405 : tensor<1x3xf64>
    %3407 = stablehlo.slice %3231 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %3408 = stablehlo.reshape %3407 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3409 = stablehlo.slice %3232 [0:1, 0:1] : (tensor<1x1xf64>) -> tensor<1x1xf64>
    %3410 = stablehlo.reshape %3409 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3411 = stablehlo.subtract %3408, %3410 : tensor<1xf64>
    %3412 = stablehlo.slice %3230 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %3413 = stablehlo.reshape %3412 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %cst_301 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %3414 = stablehlo.broadcast_in_dim %cst_301, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3415 = stablehlo.multiply %3411, %3414 : tensor<1xf64>
    %3416 = stablehlo.sine %3415 : tensor<1xf64>
    %cst_302 = stablehlo.constant dense<5.000000e-01> : tensor<f64>
    %3417 = stablehlo.broadcast_in_dim %cst_302, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3418 = stablehlo.multiply %3411, %3417 : tensor<1xf64>
    %3419 = stablehlo.cosine %3418 : tensor<1xf64>
    %3420 = stablehlo.broadcast_in_dim %3416, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3421 = stablehlo.broadcast_in_dim %3420, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3422 = stablehlo.multiply %3413, %3421 : tensor<1x3xf64>
    %3423 = stablehlo.broadcast_in_dim %3419, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %c_303 = stablehlo.constant dense<0> : tensor<i64>
    %3424 = stablehlo.reshape %c_303 : (tensor<i64>) -> tensor<1xi64>
    %3425 = stablehlo.slice %3424 [0:1] : (tensor<1xi64>) -> tensor<1xi64>
    %3426 = stablehlo.reshape %3425 : (tensor<1xi64>) -> tensor<i64>
    %3427 = stablehlo.broadcast_in_dim %3426, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %cst_304 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %3428 = stablehlo.broadcast_in_dim %cst_304, dims = [] : (tensor<f64>) -> tensor<4xf64>
    %c_305 = stablehlo.constant dense<0> : tensor<i64>
    %3429 = stablehlo.broadcast_in_dim %c_305, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %3430 = stablehlo.compare  LT, %3427, %3429,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_306 = stablehlo.constant dense<3> : tensor<i64>
    %3431 = stablehlo.broadcast_in_dim %c_306, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %3432 = stablehlo.add %3427, %3431 : tensor<1xi64>
    %3433 = call @_where_7(%3430, %3432, %3427) : (tensor<1xi1>, tensor<1xi64>, tensor<1xi64>) -> tensor<1xi64>
    %c_307 = stablehlo.constant dense<0> : tensor<i64>
    %c_308 = stablehlo.constant dense<3> : tensor<i64>
    %3434 = call @clip(%3433, %c_307, %c_308) : (tensor<1xi64>, tensor<i64>, tensor<i64>) -> tensor<1xi64>
    %3435 = call @argsort(%3434) : (tensor<1xi64>) -> tensor<1xi64>
    %3436 = stablehlo.iota dim = 0 : tensor<1xi64>
    %c_309 = stablehlo.constant dense<0> : tensor<i64>
    %3437 = stablehlo.broadcast_in_dim %c_309, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %3438 = stablehlo.compare  LT, %3435, %3437,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_310 = stablehlo.constant dense<1> : tensor<i64>
    %3439 = stablehlo.broadcast_in_dim %c_310, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %3440 = stablehlo.add %3435, %3439 : tensor<1xi64>
    %3441 = stablehlo.select %3438, %3440, %3435 : tensor<1xi1>, tensor<1xi64>
    %3442 = stablehlo.convert %3441 : (tensor<1xi64>) -> tensor<1xi32>
    %3443 = stablehlo.broadcast_in_dim %3442, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %3444 = stablehlo.convert %3434 : tensor<1xi64>
    %3445 = "stablehlo.scatter"(%3444, %3443, %3436) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<1xi64>, tensor<1x1xi32>, tensor<1xi64>) -> tensor<1xi64>
    %3446 = stablehlo.convert %3445 : tensor<1xi64>
    %c_311 = stablehlo.constant dense<true> : tensor<i1>
    %3447 = stablehlo.broadcast_in_dim %c_311, dims = [] : (tensor<i1>) -> tensor<4xi1>
    %c_312 = stablehlo.constant dense<0> : tensor<i64>
    %3448 = stablehlo.broadcast_in_dim %c_312, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %3449 = stablehlo.compare  LT, %3446, %3448,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_313 = stablehlo.constant dense<4> : tensor<i64>
    %3450 = stablehlo.broadcast_in_dim %c_313, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %3451 = stablehlo.add %3446, %3450 : tensor<1xi64>
    %3452 = stablehlo.select %3449, %3451, %3446 : tensor<1xi1>, tensor<1xi64>
    %3453 = stablehlo.convert %3452 : (tensor<1xi64>) -> tensor<1xi32>
    %3454 = stablehlo.broadcast_in_dim %3453, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %c_314 = stablehlo.constant dense<false> : tensor<i1>
    %3455 = stablehlo.broadcast_in_dim %c_314, dims = [] : (tensor<i1>) -> tensor<1xi1>
    %3456 = "stablehlo.scatter"(%3447, %3454, %3455) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i1>, %arg114: tensor<i1>):
      stablehlo.return %arg114 : tensor<i1>
    }) : (tensor<4xi1>, tensor<1x1xi32>, tensor<1xi1>) -> tensor<4xi1>
    %3457 = call @_cumulative_reduction(%3456) : (tensor<4xi1>) -> tensor<4xi64>
    %c_315 = stablehlo.constant dense<0> : tensor<i64>
    %3458 = stablehlo.broadcast_in_dim %c_315, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %c_316 = stablehlo.constant dense<0> : tensor<i64>
    %3459 = call @clip_8(%3457, %c_316) : (tensor<4xi64>, tensor<i64>) -> tensor<4xi64>
    %c_317 = stablehlo.constant dense<1> : tensor<i64>
    %c_318 = stablehlo.constant dense<0> : tensor<i64>
    %3460 = stablehlo.broadcast_in_dim %c_318, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %3461 = stablehlo.compare  LT, %3459, %3460,  SIGNED : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi1>
    %c_319 = stablehlo.constant dense<3> : tensor<i64>
    %3462 = stablehlo.broadcast_in_dim %c_319, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %3463 = stablehlo.add %3459, %3462 : tensor<4xi64>
    %3464 = stablehlo.select %3461, %3463, %3459 : tensor<4xi1>, tensor<4xi64>
    %3465 = stablehlo.convert %3464 : (tensor<4xi64>) -> tensor<4xi32>
    %3466 = stablehlo.broadcast_in_dim %3465, dims = [0] : (tensor<4xi32>) -> tensor<4x1xi32>
    %3467 = stablehlo.broadcast_in_dim %c_317, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %3468 = "stablehlo.scatter"(%3458, %3466, %3467) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0], scatter_dims_to_operand_dims = [0], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<i64>, %arg114: tensor<i64>):
      %4428 = stablehlo.add %arg113, %arg114 : tensor<i64>
      stablehlo.return %4428 : tensor<i64>
    }) : (tensor<3xi64>, tensor<4x1xi32>, tensor<4xi64>) -> tensor<3xi64>
    %3469 = call @_cumulative_reduction_9(%3468) : (tensor<3xi64>) -> tensor<3xi64>
    %c_320 = stablehlo.constant dense<1> : tensor<i64>
    %3470 = call @floor_divide(%3469, %c_320) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_321 = stablehlo.constant dense<4> : tensor<i64>
    %3471 = call @remainder_12(%3470, %c_321) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    %c_322 = stablehlo.constant dense<0> : tensor<i64>
    %3472 = stablehlo.broadcast_in_dim %c_322, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %3473 = stablehlo.compare  LT, %3446, %3472,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_323 = stablehlo.constant dense<4> : tensor<i64>
    %3474 = stablehlo.broadcast_in_dim %c_323, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %3475 = stablehlo.add %3446, %3474 : tensor<1xi64>
    %3476 = stablehlo.select %3473, %3475, %3446 : tensor<1xi1>, tensor<1xi64>
    %3477 = stablehlo.convert %3476 : (tensor<1xi64>) -> tensor<1xi32>
    %3478 = stablehlo.broadcast_in_dim %3477, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %3479 = stablehlo.broadcast_in_dim %3428, dims = [1] : (tensor<4xf64>) -> tensor<1x4xf64>
    %3480 = "stablehlo.scatter"(%3479, %3478, %3423) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<1x1xi32>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %c_324 = stablehlo.constant dense<0> : tensor<i64>
    %3481 = stablehlo.broadcast_in_dim %c_324, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %3482 = stablehlo.compare  LT, %3471, %3481,  SIGNED : (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi1>
    %c_325 = stablehlo.constant dense<4> : tensor<i64>
    %3483 = stablehlo.broadcast_in_dim %c_325, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %3484 = stablehlo.add %3471, %3483 : tensor<3xi64>
    %3485 = stablehlo.select %3482, %3484, %3471 : tensor<3xi1>, tensor<3xi64>
    %3486 = stablehlo.convert %3485 : (tensor<3xi64>) -> tensor<3xi32>
    %3487 = stablehlo.broadcast_in_dim %3486, dims = [0] : (tensor<3xi32>) -> tensor<3x1xi32>
    %3488 = "stablehlo.scatter"(%3480, %3487, %3422) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [0], inserted_window_dims = [1], scatter_dims_to_operand_dims = [1], index_vector_dim = 1>, unique_indices = false}> ({
    ^bb0(%arg113: tensor<f64>, %arg114: tensor<f64>):
      stablehlo.return %arg114 : tensor<f64>
    }) : (tensor<1x4xf64>, tensor<3x1xi32>, tensor<1x3xf64>) -> tensor<1x4xf64>
    %3489 = stablehlo.slice %3355 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3490 = stablehlo.reshape %3489 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3491 = stablehlo.slice %3488 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3492 = stablehlo.reshape %3491 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3493 = stablehlo.multiply %3490, %3492 : tensor<1xf64>
    %3494 = stablehlo.slice %3355 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3495 = stablehlo.reshape %3494 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3496 = stablehlo.slice %3488 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3497 = stablehlo.reshape %3496 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3498 = stablehlo.multiply %3495, %3497 : tensor<1xf64>
    %3499 = stablehlo.subtract %3493, %3498 : tensor<1xf64>
    %3500 = stablehlo.slice %3355 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3501 = stablehlo.reshape %3500 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3502 = stablehlo.slice %3488 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3503 = stablehlo.reshape %3502 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3504 = stablehlo.multiply %3501, %3503 : tensor<1xf64>
    %3505 = stablehlo.subtract %3499, %3504 : tensor<1xf64>
    %3506 = stablehlo.slice %3355 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3507 = stablehlo.reshape %3506 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3508 = stablehlo.slice %3488 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3509 = stablehlo.reshape %3508 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3510 = stablehlo.multiply %3507, %3509 : tensor<1xf64>
    %3511 = stablehlo.subtract %3505, %3510 : tensor<1xf64>
    %3512 = stablehlo.slice %3355 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3513 = stablehlo.reshape %3512 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3514 = stablehlo.slice %3488 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3515 = stablehlo.reshape %3514 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3516 = stablehlo.multiply %3513, %3515 : tensor<1xf64>
    %3517 = stablehlo.slice %3355 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3518 = stablehlo.reshape %3517 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3519 = stablehlo.slice %3488 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3520 = stablehlo.reshape %3519 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3521 = stablehlo.multiply %3518, %3520 : tensor<1xf64>
    %3522 = stablehlo.add %3516, %3521 : tensor<1xf64>
    %3523 = stablehlo.slice %3355 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3524 = stablehlo.reshape %3523 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3525 = stablehlo.slice %3488 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3526 = stablehlo.reshape %3525 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3527 = stablehlo.multiply %3524, %3526 : tensor<1xf64>
    %3528 = stablehlo.add %3522, %3527 : tensor<1xf64>
    %3529 = stablehlo.slice %3355 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3530 = stablehlo.reshape %3529 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3531 = stablehlo.slice %3488 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3532 = stablehlo.reshape %3531 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3533 = stablehlo.multiply %3530, %3532 : tensor<1xf64>
    %3534 = stablehlo.subtract %3528, %3533 : tensor<1xf64>
    %3535 = stablehlo.slice %3355 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3536 = stablehlo.reshape %3535 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3537 = stablehlo.slice %3488 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3538 = stablehlo.reshape %3537 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3539 = stablehlo.multiply %3536, %3538 : tensor<1xf64>
    %3540 = stablehlo.slice %3355 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3541 = stablehlo.reshape %3540 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3542 = stablehlo.slice %3488 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3543 = stablehlo.reshape %3542 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3544 = stablehlo.multiply %3541, %3543 : tensor<1xf64>
    %3545 = stablehlo.subtract %3539, %3544 : tensor<1xf64>
    %3546 = stablehlo.slice %3355 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3547 = stablehlo.reshape %3546 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3548 = stablehlo.slice %3488 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3549 = stablehlo.reshape %3548 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3550 = stablehlo.multiply %3547, %3549 : tensor<1xf64>
    %3551 = stablehlo.add %3545, %3550 : tensor<1xf64>
    %3552 = stablehlo.slice %3355 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3553 = stablehlo.reshape %3552 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3554 = stablehlo.slice %3488 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3555 = stablehlo.reshape %3554 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3556 = stablehlo.multiply %3553, %3555 : tensor<1xf64>
    %3557 = stablehlo.add %3551, %3556 : tensor<1xf64>
    %3558 = stablehlo.slice %3355 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3559 = stablehlo.reshape %3558 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3560 = stablehlo.slice %3488 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3561 = stablehlo.reshape %3560 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3562 = stablehlo.multiply %3559, %3561 : tensor<1xf64>
    %3563 = stablehlo.slice %3355 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3564 = stablehlo.reshape %3563 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3565 = stablehlo.slice %3488 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3566 = stablehlo.reshape %3565 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3567 = stablehlo.multiply %3564, %3566 : tensor<1xf64>
    %3568 = stablehlo.add %3562, %3567 : tensor<1xf64>
    %3569 = stablehlo.slice %3355 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3570 = stablehlo.reshape %3569 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3571 = stablehlo.slice %3488 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3572 = stablehlo.reshape %3571 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3573 = stablehlo.multiply %3570, %3572 : tensor<1xf64>
    %3574 = stablehlo.subtract %3568, %3573 : tensor<1xf64>
    %3575 = stablehlo.slice %3355 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3576 = stablehlo.reshape %3575 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3577 = stablehlo.slice %3488 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3578 = stablehlo.reshape %3577 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3579 = stablehlo.multiply %3576, %3578 : tensor<1xf64>
    %3580 = stablehlo.add %3574, %3579 : tensor<1xf64>
    %3581 = stablehlo.broadcast_in_dim %3511, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3582 = stablehlo.broadcast_in_dim %3534, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3583 = stablehlo.broadcast_in_dim %3557, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3584 = stablehlo.broadcast_in_dim %3580, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3585 = stablehlo.concatenate %3581, %3582, %3583, %3584, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %3586 = stablehlo.slice %3229 [0:1, 0:1, 0:3] : (tensor<1x1x3xf64>) -> tensor<1x1x3xf64>
    %3587 = stablehlo.reshape %3586 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3588 = stablehlo.slice %3585 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %3589 = stablehlo.reshape %3588 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3590 = stablehlo.slice %3585 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %3591 = stablehlo.dot_general %3590, %3587, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %3592 = stablehlo.broadcast_in_dim %3591, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3593 = stablehlo.broadcast_in_dim %3592, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3594 = stablehlo.multiply %3593, %3590 : tensor<1x3xf64>
    %cst_326 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3595 = stablehlo.broadcast_in_dim %cst_326, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %3596 = stablehlo.multiply %3595, %3594 : tensor<1x3xf64>
    %3597 = stablehlo.multiply %3589, %3589 : tensor<1xf64>
    %3598 = stablehlo.dot_general %3590, %3590, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %3599 = stablehlo.subtract %3597, %3598 : tensor<1xf64>
    %3600 = stablehlo.broadcast_in_dim %3599, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3601 = stablehlo.broadcast_in_dim %3600, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3602 = stablehlo.multiply %3601, %3587 : tensor<1x3xf64>
    %3603 = stablehlo.add %3596, %3602 : tensor<1x3xf64>
    %cst_327 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3604 = stablehlo.broadcast_in_dim %cst_327, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3605 = stablehlo.multiply %3604, %3589 : tensor<1xf64>
    %3606 = call @cross(%3590, %3587) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %3607 = stablehlo.broadcast_in_dim %3605, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3608 = stablehlo.broadcast_in_dim %3607, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %3609 = stablehlo.multiply %3608, %3606 : tensor<1x3xf64>
    %3610 = stablehlo.add %3603, %3609 : tensor<1x3xf64>
    %3611 = stablehlo.subtract %3381, %3610 : tensor<1x3xf64>
    %3612 = stablehlo.broadcast_in_dim %3381, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %3613 = stablehlo.broadcast_in_dim %3406, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %3614 = stablehlo.broadcast_in_dim %3585, dims = [0, 1] : (tensor<1x4xf64>) -> tensor<1x4x1xf64>
    %3615 = stablehlo.broadcast_in_dim %3585, dims = [0, 2] : (tensor<1x4xf64>) -> tensor<1x1x4xf64>
    %3616 = stablehlo.broadcast_in_dim %3614, dims = [0, 1, 2] : (tensor<1x4x1xf64>) -> tensor<1x4x4xf64>
    %3617 = stablehlo.broadcast_in_dim %3615, dims = [0, 1, 2] : (tensor<1x1x4xf64>) -> tensor<1x4x4xf64>
    %3618 = stablehlo.multiply %3616, %3617 : tensor<1x4x4xf64>
    %3619 = stablehlo.slice %3618 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3620 = stablehlo.reshape %3619 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3621 = stablehlo.slice %3618 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3622 = stablehlo.reshape %3621 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3623 = stablehlo.add %3620, %3622 : tensor<1xf64>
    %3624 = stablehlo.slice %3618 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3625 = stablehlo.reshape %3624 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3626 = stablehlo.subtract %3623, %3625 : tensor<1xf64>
    %3627 = stablehlo.slice %3618 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3628 = stablehlo.reshape %3627 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3629 = stablehlo.subtract %3626, %3628 : tensor<1xf64>
    %3630 = stablehlo.slice %3618 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3631 = stablehlo.reshape %3630 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3632 = stablehlo.slice %3618 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3633 = stablehlo.reshape %3632 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3634 = stablehlo.subtract %3631, %3633 : tensor<1xf64>
    %cst_328 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3635 = stablehlo.broadcast_in_dim %cst_328, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3636 = stablehlo.multiply %3635, %3634 : tensor<1xf64>
    %3637 = stablehlo.slice %3618 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3638 = stablehlo.reshape %3637 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3639 = stablehlo.slice %3618 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3640 = stablehlo.reshape %3639 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3641 = stablehlo.add %3638, %3640 : tensor<1xf64>
    %cst_329 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3642 = stablehlo.broadcast_in_dim %cst_329, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3643 = stablehlo.multiply %3642, %3641 : tensor<1xf64>
    %3644 = stablehlo.slice %3618 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3645 = stablehlo.reshape %3644 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3646 = stablehlo.slice %3618 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3647 = stablehlo.reshape %3646 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3648 = stablehlo.add %3645, %3647 : tensor<1xf64>
    %cst_330 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3649 = stablehlo.broadcast_in_dim %cst_330, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3650 = stablehlo.multiply %3649, %3648 : tensor<1xf64>
    %3651 = stablehlo.slice %3618 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3652 = stablehlo.reshape %3651 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3653 = stablehlo.slice %3618 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3654 = stablehlo.reshape %3653 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3655 = stablehlo.subtract %3652, %3654 : tensor<1xf64>
    %3656 = stablehlo.slice %3618 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3657 = stablehlo.reshape %3656 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3658 = stablehlo.add %3655, %3657 : tensor<1xf64>
    %3659 = stablehlo.slice %3618 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3660 = stablehlo.reshape %3659 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3661 = stablehlo.subtract %3658, %3660 : tensor<1xf64>
    %3662 = stablehlo.slice %3618 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3663 = stablehlo.reshape %3662 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3664 = stablehlo.slice %3618 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3665 = stablehlo.reshape %3664 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3666 = stablehlo.subtract %3663, %3665 : tensor<1xf64>
    %cst_331 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3667 = stablehlo.broadcast_in_dim %cst_331, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3668 = stablehlo.multiply %3667, %3666 : tensor<1xf64>
    %3669 = stablehlo.slice %3618 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3670 = stablehlo.reshape %3669 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3671 = stablehlo.slice %3618 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3672 = stablehlo.reshape %3671 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3673 = stablehlo.subtract %3670, %3672 : tensor<1xf64>
    %cst_332 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3674 = stablehlo.broadcast_in_dim %cst_332, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3675 = stablehlo.multiply %3674, %3673 : tensor<1xf64>
    %3676 = stablehlo.slice %3618 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3677 = stablehlo.reshape %3676 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3678 = stablehlo.slice %3618 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3679 = stablehlo.reshape %3678 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3680 = stablehlo.add %3677, %3679 : tensor<1xf64>
    %cst_333 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3681 = stablehlo.broadcast_in_dim %cst_333, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %3682 = stablehlo.multiply %3681, %3680 : tensor<1xf64>
    %3683 = stablehlo.slice %3618 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3684 = stablehlo.reshape %3683 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3685 = stablehlo.slice %3618 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3686 = stablehlo.reshape %3685 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3687 = stablehlo.subtract %3684, %3686 : tensor<1xf64>
    %3688 = stablehlo.slice %3618 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3689 = stablehlo.reshape %3688 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3690 = stablehlo.subtract %3687, %3689 : tensor<1xf64>
    %3691 = stablehlo.slice %3618 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %3692 = stablehlo.reshape %3691 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %3693 = stablehlo.add %3690, %3692 : tensor<1xf64>
    %3694 = stablehlo.broadcast_in_dim %3629, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3695 = stablehlo.broadcast_in_dim %3636, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3696 = stablehlo.broadcast_in_dim %3643, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3697 = stablehlo.concatenate %3694, %3695, %3696, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %3698 = stablehlo.broadcast_in_dim %3650, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3699 = stablehlo.broadcast_in_dim %3661, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3700 = stablehlo.broadcast_in_dim %3668, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3701 = stablehlo.concatenate %3698, %3699, %3700, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %3702 = stablehlo.broadcast_in_dim %3675, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3703 = stablehlo.broadcast_in_dim %3682, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3704 = stablehlo.broadcast_in_dim %3693, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %3705 = stablehlo.concatenate %3702, %3703, %3704, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %3706 = stablehlo.broadcast_in_dim %3697, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %3707 = stablehlo.broadcast_in_dim %3701, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %3708 = stablehlo.broadcast_in_dim %3705, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %3709 = stablehlo.concatenate %3706, %3707, %3708, dim = 1 : (tensor<1x1x3xf64>, tensor<1x1x3xf64>, tensor<1x1x3xf64>) -> tensor<1x3x3xf64>
    %3710 = stablehlo.reshape %0 : (tensor<1x0xf64>) -> tensor<0xf64>
    %3711 = stablehlo.reshape %105 : (tensor<1x0xf64>) -> tensor<0xf64>
    %3712 = stablehlo.reshape %333 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3713 = stablehlo.reshape %816 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3714 = stablehlo.reshape %1299 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3715 = stablehlo.reshape %1782 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3716 = stablehlo.reshape %2265 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3717 = stablehlo.reshape %2748 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3718 = stablehlo.reshape %3231 : (tensor<1x1xf64>) -> tensor<1xf64>
    %3719 = stablehlo.concatenate %3710, %3711, %3712, %3713, %3714, %3715, %3716, %3717, %3718, dim = 0 : (tensor<0xf64>, tensor<0xf64>, tensor<1xf64>, tensor<1xf64>, tensor<1xf64>, tensor<1xf64>, tensor<1xf64>, tensor<1xf64>, tensor<1xf64>) -> tensor<7xf64>
    %3720 = call @_take_13(%3719, %c_36) : (tensor<7xf64>, tensor<7xi64>) -> tensor<7xf64>
    %3721 = stablehlo.reshape %101 : (tensor<1x0x3xf64>) -> tensor<0x3xf64>
    %3722 = stablehlo.reshape %327 : (tensor<1x0x3xf64>) -> tensor<0x3xf64>
    %3723 = stablehlo.reshape %714 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3724 = stablehlo.reshape %1197 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3725 = stablehlo.reshape %1680 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3726 = stablehlo.reshape %2163 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3727 = stablehlo.reshape %2646 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3728 = stablehlo.reshape %3129 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3729 = stablehlo.reshape %3612 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3730 = stablehlo.concatenate %3721, %3722, %3723, %3724, %3725, %3726, %3727, %3728, %3729, dim = 0 : (tensor<0x3xf64>, tensor<0x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<7x3xf64>
    %3731 = call @_take_15(%3730, %c_37) : (tensor<7x3xf64>, tensor<7xi64>) -> tensor<7x3xf64>
    %3732 = stablehlo.reshape %102 : (tensor<1x0x3xf64>) -> tensor<0x3xf64>
    %3733 = stablehlo.reshape %328 : (tensor<1x0x3xf64>) -> tensor<0x3xf64>
    %3734 = stablehlo.reshape %715 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3735 = stablehlo.reshape %1198 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3736 = stablehlo.reshape %1681 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3737 = stablehlo.reshape %2164 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3738 = stablehlo.reshape %2647 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3739 = stablehlo.reshape %3130 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3740 = stablehlo.reshape %3613 : (tensor<1x1x3xf64>) -> tensor<1x3xf64>
    %3741 = stablehlo.concatenate %3732, %3733, %3734, %3735, %3736, %3737, %3738, %3739, %3740, dim = 0 : (tensor<0x3xf64>, tensor<0x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<7x3xf64>
    %3742 = call @_take_15(%3741, %c_38) : (tensor<7x3xf64>, tensor<7xi64>) -> tensor<7x3xf64>
    %3743 = stablehlo.concatenate %1, %131, %713, %1196, %1679, %2162, %2645, %3128, %3611, dim = 0 : (tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<9x3xf64>
    %3744 = call @_take_16(%3743, %c_39) : (tensor<9x3xf64>, tensor<9xi64>) -> tensor<9x3xf64>
    %3745 = stablehlo.concatenate %2, %228, %687, %1170, %1653, %2136, %2619, %3102, %3585, dim = 0 : (tensor<1x4xf64>, tensor<1x4xf64>, tensor<1x4xf64>, tensor<1x4xf64>, tensor<1x4xf64>, tensor<1x4xf64>, tensor<1x4xf64>, tensor<1x4xf64>, tensor<1x4xf64>) -> tensor<9x4xf64>
    %3746 = call @_take_18(%3745, %c_40) : (tensor<9x4xf64>, tensor<9xi64>) -> tensor<9x4xf64>
    %3747 = stablehlo.concatenate %100, %326, %811, %1294, %1777, %2260, %2743, %3226, %3709, dim = 0 : (tensor<1x3x3xf64>, tensor<1x3x3xf64>, tensor<1x3x3xf64>, tensor<1x3x3xf64>, tensor<1x3x3xf64>, tensor<1x3x3xf64>, tensor<1x3x3xf64>, tensor<1x3x3xf64>, tensor<1x3x3xf64>) -> tensor<9x3x3xf64>
    %3748 = call @_take_19(%3747, %c_41) : (tensor<9x3x3xf64>, tensor<9xi64>) -> tensor<9x3x3xf64>
    %3749 = stablehlo.slice %3746 [0:9, 0:1] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3750 = stablehlo.reshape %3749 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3751 = stablehlo.slice %3746 [0:9, 1:4] : (tensor<9x4xf64>) -> tensor<9x3xf64>
    %3752 = stablehlo.dot_general %3751, %arg3, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<9x3xf64>, tensor<9x3xf64>) -> tensor<9xf64>
    %3753 = stablehlo.broadcast_in_dim %3752, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3754 = stablehlo.broadcast_in_dim %3753, dims = [0, 1] : (tensor<9x1xf64>) -> tensor<9x3xf64>
    %3755 = stablehlo.multiply %3754, %3751 : tensor<9x3xf64>
    %cst_334 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3756 = stablehlo.broadcast_in_dim %cst_334, dims = [] : (tensor<f64>) -> tensor<9x3xf64>
    %3757 = stablehlo.multiply %3756, %3755 : tensor<9x3xf64>
    %3758 = stablehlo.multiply %3750, %3750 : tensor<9xf64>
    %3759 = stablehlo.dot_general %3751, %3751, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<9x3xf64>, tensor<9x3xf64>) -> tensor<9xf64>
    %3760 = stablehlo.subtract %3758, %3759 : tensor<9xf64>
    %3761 = stablehlo.broadcast_in_dim %3760, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3762 = stablehlo.broadcast_in_dim %3761, dims = [0, 1] : (tensor<9x1xf64>) -> tensor<9x3xf64>
    %3763 = stablehlo.multiply %3762, %arg3 : tensor<9x3xf64>
    %3764 = stablehlo.add %3757, %3763 : tensor<9x3xf64>
    %cst_335 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3765 = stablehlo.broadcast_in_dim %cst_335, dims = [] : (tensor<f64>) -> tensor<9xf64>
    %3766 = stablehlo.multiply %3765, %3750 : tensor<9xf64>
    %3767 = call @cross_20(%3751, %arg3) : (tensor<9x3xf64>, tensor<9x3xf64>) -> tensor<9x3xf64>
    %3768 = stablehlo.broadcast_in_dim %3766, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3769 = stablehlo.broadcast_in_dim %3768, dims = [0, 1] : (tensor<9x1xf64>) -> tensor<9x3xf64>
    %3770 = stablehlo.multiply %3769, %3767 : tensor<9x3xf64>
    %3771 = stablehlo.add %3764, %3770 : tensor<9x3xf64>
    %3772 = stablehlo.add %3744, %3771 : tensor<9x3xf64>
    %3773 = stablehlo.slice %3746 [0:9, 0:1] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3774 = stablehlo.reshape %3773 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3775 = stablehlo.slice %arg4 [0:9, 0:1] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3776 = stablehlo.reshape %3775 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3777 = stablehlo.multiply %3774, %3776 : tensor<9xf64>
    %3778 = stablehlo.slice %3746 [0:9, 1:2] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3779 = stablehlo.reshape %3778 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3780 = stablehlo.slice %arg4 [0:9, 1:2] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3781 = stablehlo.reshape %3780 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3782 = stablehlo.multiply %3779, %3781 : tensor<9xf64>
    %3783 = stablehlo.subtract %3777, %3782 : tensor<9xf64>
    %3784 = stablehlo.slice %3746 [0:9, 2:3] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3785 = stablehlo.reshape %3784 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3786 = stablehlo.slice %arg4 [0:9, 2:3] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3787 = stablehlo.reshape %3786 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3788 = stablehlo.multiply %3785, %3787 : tensor<9xf64>
    %3789 = stablehlo.subtract %3783, %3788 : tensor<9xf64>
    %3790 = stablehlo.slice %3746 [0:9, 3:4] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3791 = stablehlo.reshape %3790 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3792 = stablehlo.slice %arg4 [0:9, 3:4] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3793 = stablehlo.reshape %3792 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3794 = stablehlo.multiply %3791, %3793 : tensor<9xf64>
    %3795 = stablehlo.subtract %3789, %3794 : tensor<9xf64>
    %3796 = stablehlo.slice %3746 [0:9, 0:1] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3797 = stablehlo.reshape %3796 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3798 = stablehlo.slice %arg4 [0:9, 1:2] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3799 = stablehlo.reshape %3798 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3800 = stablehlo.multiply %3797, %3799 : tensor<9xf64>
    %3801 = stablehlo.slice %3746 [0:9, 1:2] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3802 = stablehlo.reshape %3801 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3803 = stablehlo.slice %arg4 [0:9, 0:1] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3804 = stablehlo.reshape %3803 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3805 = stablehlo.multiply %3802, %3804 : tensor<9xf64>
    %3806 = stablehlo.add %3800, %3805 : tensor<9xf64>
    %3807 = stablehlo.slice %3746 [0:9, 2:3] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3808 = stablehlo.reshape %3807 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3809 = stablehlo.slice %arg4 [0:9, 3:4] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3810 = stablehlo.reshape %3809 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3811 = stablehlo.multiply %3808, %3810 : tensor<9xf64>
    %3812 = stablehlo.add %3806, %3811 : tensor<9xf64>
    %3813 = stablehlo.slice %3746 [0:9, 3:4] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3814 = stablehlo.reshape %3813 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3815 = stablehlo.slice %arg4 [0:9, 2:3] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3816 = stablehlo.reshape %3815 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3817 = stablehlo.multiply %3814, %3816 : tensor<9xf64>
    %3818 = stablehlo.subtract %3812, %3817 : tensor<9xf64>
    %3819 = stablehlo.slice %3746 [0:9, 0:1] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3820 = stablehlo.reshape %3819 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3821 = stablehlo.slice %arg4 [0:9, 2:3] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3822 = stablehlo.reshape %3821 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3823 = stablehlo.multiply %3820, %3822 : tensor<9xf64>
    %3824 = stablehlo.slice %3746 [0:9, 1:2] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3825 = stablehlo.reshape %3824 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3826 = stablehlo.slice %arg4 [0:9, 3:4] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3827 = stablehlo.reshape %3826 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3828 = stablehlo.multiply %3825, %3827 : tensor<9xf64>
    %3829 = stablehlo.subtract %3823, %3828 : tensor<9xf64>
    %3830 = stablehlo.slice %3746 [0:9, 2:3] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3831 = stablehlo.reshape %3830 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3832 = stablehlo.slice %arg4 [0:9, 0:1] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3833 = stablehlo.reshape %3832 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3834 = stablehlo.multiply %3831, %3833 : tensor<9xf64>
    %3835 = stablehlo.add %3829, %3834 : tensor<9xf64>
    %3836 = stablehlo.slice %3746 [0:9, 3:4] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3837 = stablehlo.reshape %3836 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3838 = stablehlo.slice %arg4 [0:9, 1:2] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3839 = stablehlo.reshape %3838 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3840 = stablehlo.multiply %3837, %3839 : tensor<9xf64>
    %3841 = stablehlo.add %3835, %3840 : tensor<9xf64>
    %3842 = stablehlo.slice %3746 [0:9, 0:1] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3843 = stablehlo.reshape %3842 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3844 = stablehlo.slice %arg4 [0:9, 3:4] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3845 = stablehlo.reshape %3844 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3846 = stablehlo.multiply %3843, %3845 : tensor<9xf64>
    %3847 = stablehlo.slice %3746 [0:9, 1:2] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3848 = stablehlo.reshape %3847 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3849 = stablehlo.slice %arg4 [0:9, 2:3] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3850 = stablehlo.reshape %3849 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3851 = stablehlo.multiply %3848, %3850 : tensor<9xf64>
    %3852 = stablehlo.add %3846, %3851 : tensor<9xf64>
    %3853 = stablehlo.slice %3746 [0:9, 2:3] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3854 = stablehlo.reshape %3853 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3855 = stablehlo.slice %arg4 [0:9, 1:2] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3856 = stablehlo.reshape %3855 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3857 = stablehlo.multiply %3854, %3856 : tensor<9xf64>
    %3858 = stablehlo.subtract %3852, %3857 : tensor<9xf64>
    %3859 = stablehlo.slice %3746 [0:9, 3:4] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3860 = stablehlo.reshape %3859 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3861 = stablehlo.slice %arg4 [0:9, 0:1] : (tensor<9x4xf64>) -> tensor<9x1xf64>
    %3862 = stablehlo.reshape %3861 : (tensor<9x1xf64>) -> tensor<9xf64>
    %3863 = stablehlo.multiply %3860, %3862 : tensor<9xf64>
    %3864 = stablehlo.add %3858, %3863 : tensor<9xf64>
    %3865 = stablehlo.broadcast_in_dim %3795, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3866 = stablehlo.broadcast_in_dim %3818, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3867 = stablehlo.broadcast_in_dim %3841, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3868 = stablehlo.broadcast_in_dim %3864, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3869 = stablehlo.concatenate %3865, %3866, %3867, %3868, dim = 1 : (tensor<9x1xf64>, tensor<9x1xf64>, tensor<9x1xf64>, tensor<9x1xf64>) -> tensor<9x4xf64>
    %3870 = stablehlo.broadcast_in_dim %3869, dims = [0, 1] : (tensor<9x4xf64>) -> tensor<9x4x1xf64>
    %3871 = stablehlo.broadcast_in_dim %3869, dims = [0, 2] : (tensor<9x4xf64>) -> tensor<9x1x4xf64>
    %3872 = stablehlo.broadcast_in_dim %3870, dims = [0, 1, 2] : (tensor<9x4x1xf64>) -> tensor<9x4x4xf64>
    %3873 = stablehlo.broadcast_in_dim %3871, dims = [0, 1, 2] : (tensor<9x1x4xf64>) -> tensor<9x4x4xf64>
    %3874 = stablehlo.multiply %3872, %3873 : tensor<9x4x4xf64>
    %3875 = stablehlo.slice %3874 [0:9, 0:1, 0:1] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3876 = stablehlo.reshape %3875 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3877 = stablehlo.slice %3874 [0:9, 1:2, 1:2] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3878 = stablehlo.reshape %3877 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3879 = stablehlo.add %3876, %3878 : tensor<9xf64>
    %3880 = stablehlo.slice %3874 [0:9, 2:3, 2:3] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3881 = stablehlo.reshape %3880 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3882 = stablehlo.subtract %3879, %3881 : tensor<9xf64>
    %3883 = stablehlo.slice %3874 [0:9, 3:4, 3:4] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3884 = stablehlo.reshape %3883 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3885 = stablehlo.subtract %3882, %3884 : tensor<9xf64>
    %3886 = stablehlo.slice %3874 [0:9, 1:2, 2:3] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3887 = stablehlo.reshape %3886 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3888 = stablehlo.slice %3874 [0:9, 0:1, 3:4] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3889 = stablehlo.reshape %3888 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3890 = stablehlo.subtract %3887, %3889 : tensor<9xf64>
    %cst_336 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3891 = stablehlo.broadcast_in_dim %cst_336, dims = [] : (tensor<f64>) -> tensor<9xf64>
    %3892 = stablehlo.multiply %3891, %3890 : tensor<9xf64>
    %3893 = stablehlo.slice %3874 [0:9, 1:2, 3:4] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3894 = stablehlo.reshape %3893 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3895 = stablehlo.slice %3874 [0:9, 0:1, 2:3] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3896 = stablehlo.reshape %3895 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3897 = stablehlo.add %3894, %3896 : tensor<9xf64>
    %cst_337 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3898 = stablehlo.broadcast_in_dim %cst_337, dims = [] : (tensor<f64>) -> tensor<9xf64>
    %3899 = stablehlo.multiply %3898, %3897 : tensor<9xf64>
    %3900 = stablehlo.slice %3874 [0:9, 1:2, 2:3] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3901 = stablehlo.reshape %3900 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3902 = stablehlo.slice %3874 [0:9, 0:1, 3:4] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3903 = stablehlo.reshape %3902 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3904 = stablehlo.add %3901, %3903 : tensor<9xf64>
    %cst_338 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3905 = stablehlo.broadcast_in_dim %cst_338, dims = [] : (tensor<f64>) -> tensor<9xf64>
    %3906 = stablehlo.multiply %3905, %3904 : tensor<9xf64>
    %3907 = stablehlo.slice %3874 [0:9, 0:1, 0:1] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3908 = stablehlo.reshape %3907 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3909 = stablehlo.slice %3874 [0:9, 1:2, 1:2] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3910 = stablehlo.reshape %3909 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3911 = stablehlo.subtract %3908, %3910 : tensor<9xf64>
    %3912 = stablehlo.slice %3874 [0:9, 2:3, 2:3] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3913 = stablehlo.reshape %3912 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3914 = stablehlo.add %3911, %3913 : tensor<9xf64>
    %3915 = stablehlo.slice %3874 [0:9, 3:4, 3:4] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3916 = stablehlo.reshape %3915 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3917 = stablehlo.subtract %3914, %3916 : tensor<9xf64>
    %3918 = stablehlo.slice %3874 [0:9, 2:3, 3:4] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3919 = stablehlo.reshape %3918 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3920 = stablehlo.slice %3874 [0:9, 0:1, 1:2] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3921 = stablehlo.reshape %3920 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3922 = stablehlo.subtract %3919, %3921 : tensor<9xf64>
    %cst_339 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3923 = stablehlo.broadcast_in_dim %cst_339, dims = [] : (tensor<f64>) -> tensor<9xf64>
    %3924 = stablehlo.multiply %3923, %3922 : tensor<9xf64>
    %3925 = stablehlo.slice %3874 [0:9, 1:2, 3:4] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3926 = stablehlo.reshape %3925 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3927 = stablehlo.slice %3874 [0:9, 0:1, 2:3] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3928 = stablehlo.reshape %3927 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3929 = stablehlo.subtract %3926, %3928 : tensor<9xf64>
    %cst_340 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3930 = stablehlo.broadcast_in_dim %cst_340, dims = [] : (tensor<f64>) -> tensor<9xf64>
    %3931 = stablehlo.multiply %3930, %3929 : tensor<9xf64>
    %3932 = stablehlo.slice %3874 [0:9, 2:3, 3:4] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3933 = stablehlo.reshape %3932 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3934 = stablehlo.slice %3874 [0:9, 0:1, 1:2] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3935 = stablehlo.reshape %3934 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3936 = stablehlo.add %3933, %3935 : tensor<9xf64>
    %cst_341 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3937 = stablehlo.broadcast_in_dim %cst_341, dims = [] : (tensor<f64>) -> tensor<9xf64>
    %3938 = stablehlo.multiply %3937, %3936 : tensor<9xf64>
    %3939 = stablehlo.slice %3874 [0:9, 0:1, 0:1] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3940 = stablehlo.reshape %3939 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3941 = stablehlo.slice %3874 [0:9, 1:2, 1:2] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3942 = stablehlo.reshape %3941 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3943 = stablehlo.subtract %3940, %3942 : tensor<9xf64>
    %3944 = stablehlo.slice %3874 [0:9, 2:3, 2:3] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3945 = stablehlo.reshape %3944 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3946 = stablehlo.subtract %3943, %3945 : tensor<9xf64>
    %3947 = stablehlo.slice %3874 [0:9, 3:4, 3:4] : (tensor<9x4x4xf64>) -> tensor<9x1x1xf64>
    %3948 = stablehlo.reshape %3947 : (tensor<9x1x1xf64>) -> tensor<9xf64>
    %3949 = stablehlo.add %3946, %3948 : tensor<9xf64>
    %3950 = stablehlo.broadcast_in_dim %3885, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3951 = stablehlo.broadcast_in_dim %3892, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3952 = stablehlo.broadcast_in_dim %3899, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3953 = stablehlo.concatenate %3950, %3951, %3952, dim = 1 : (tensor<9x1xf64>, tensor<9x1xf64>, tensor<9x1xf64>) -> tensor<9x3xf64>
    %3954 = stablehlo.broadcast_in_dim %3906, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3955 = stablehlo.broadcast_in_dim %3917, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3956 = stablehlo.broadcast_in_dim %3924, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3957 = stablehlo.concatenate %3954, %3955, %3956, dim = 1 : (tensor<9x1xf64>, tensor<9x1xf64>, tensor<9x1xf64>) -> tensor<9x3xf64>
    %3958 = stablehlo.broadcast_in_dim %3931, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3959 = stablehlo.broadcast_in_dim %3938, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3960 = stablehlo.broadcast_in_dim %3949, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %3961 = stablehlo.concatenate %3958, %3959, %3960, dim = 1 : (tensor<9x1xf64>, tensor<9x1xf64>, tensor<9x1xf64>) -> tensor<9x3xf64>
    %3962 = stablehlo.broadcast_in_dim %3953, dims = [0, 2] : (tensor<9x3xf64>) -> tensor<9x1x3xf64>
    %3963 = stablehlo.broadcast_in_dim %3957, dims = [0, 2] : (tensor<9x3xf64>) -> tensor<9x1x3xf64>
    %3964 = stablehlo.broadcast_in_dim %3961, dims = [0, 2] : (tensor<9x3xf64>) -> tensor<9x1x3xf64>
    %3965 = stablehlo.concatenate %3962, %3963, %3964, dim = 1 : (tensor<9x1x3xf64>, tensor<9x1x3xf64>, tensor<9x1x3xf64>) -> tensor<9x3x3xf64>
    %c_342 = stablehlo.constant dense<0> : tensor<i32>
    %3966 = stablehlo.broadcast_in_dim %c_342, dims = [] : (tensor<i32>) -> tensor<61xi32>
    %3967 = stablehlo.compare  LT, %c_42, %3966,  SIGNED : (tensor<61xi32>, tensor<61xi32>) -> tensor<61xi1>
    %c_343 = stablehlo.constant dense<9> : tensor<i32>
    %3968 = stablehlo.broadcast_in_dim %c_343, dims = [] : (tensor<i32>) -> tensor<61xi32>
    %3969 = stablehlo.add %c_42, %3968 : tensor<61xi32>
    %3970 = stablehlo.select %3967, %3969, %c_42 : tensor<61xi1>, tensor<61xi32>
    %3971 = stablehlo.broadcast_in_dim %3970, dims = [0] : (tensor<61xi32>) -> tensor<61x1xi32>
    %3972 = "stablehlo.gather"(%3744, %3971) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 3>}> : (tensor<9x3xf64>, tensor<61x1xi32>) -> tensor<61x3xf64>
    %c_344 = stablehlo.constant dense<0> : tensor<i32>
    %3973 = stablehlo.broadcast_in_dim %c_344, dims = [] : (tensor<i32>) -> tensor<61xi32>
    %3974 = stablehlo.compare  LT, %c_42, %3973,  SIGNED : (tensor<61xi32>, tensor<61xi32>) -> tensor<61xi1>
    %c_345 = stablehlo.constant dense<9> : tensor<i32>
    %3975 = stablehlo.broadcast_in_dim %c_345, dims = [] : (tensor<i32>) -> tensor<61xi32>
    %3976 = stablehlo.add %c_42, %3975 : tensor<61xi32>
    %3977 = stablehlo.select %3974, %3976, %c_42 : tensor<61xi1>, tensor<61xi32>
    %3978 = stablehlo.broadcast_in_dim %3977, dims = [0] : (tensor<61xi32>) -> tensor<61x1xi32>
    %3979 = "stablehlo.gather"(%3746, %3978) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 4>}> : (tensor<9x4xf64>, tensor<61x1xi32>) -> tensor<61x4xf64>
    %3980 = stablehlo.slice %3979 [0:61, 0:1] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %3981 = stablehlo.reshape %3980 : (tensor<61x1xf64>) -> tensor<61xf64>
    %3982 = stablehlo.slice %3979 [0:61, 1:4] : (tensor<61x4xf64>) -> tensor<61x3xf64>
    %3983 = stablehlo.dot_general %3982, %arg7, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<61x3xf64>, tensor<61x3xf64>) -> tensor<61xf64>
    %3984 = stablehlo.broadcast_in_dim %3983, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %3985 = stablehlo.broadcast_in_dim %3984, dims = [0, 1] : (tensor<61x1xf64>) -> tensor<61x3xf64>
    %3986 = stablehlo.multiply %3985, %3982 : tensor<61x3xf64>
    %cst_346 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3987 = stablehlo.broadcast_in_dim %cst_346, dims = [] : (tensor<f64>) -> tensor<61x3xf64>
    %3988 = stablehlo.multiply %3987, %3986 : tensor<61x3xf64>
    %3989 = stablehlo.multiply %3981, %3981 : tensor<61xf64>
    %3990 = stablehlo.dot_general %3982, %3982, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<61x3xf64>, tensor<61x3xf64>) -> tensor<61xf64>
    %3991 = stablehlo.subtract %3989, %3990 : tensor<61xf64>
    %3992 = stablehlo.broadcast_in_dim %3991, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %3993 = stablehlo.broadcast_in_dim %3992, dims = [0, 1] : (tensor<61x1xf64>) -> tensor<61x3xf64>
    %3994 = stablehlo.multiply %3993, %arg7 : tensor<61x3xf64>
    %3995 = stablehlo.add %3988, %3994 : tensor<61x3xf64>
    %cst_347 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %3996 = stablehlo.broadcast_in_dim %cst_347, dims = [] : (tensor<f64>) -> tensor<61xf64>
    %3997 = stablehlo.multiply %3996, %3981 : tensor<61xf64>
    %3998 = call @cross_21(%3982, %arg7) : (tensor<61x3xf64>, tensor<61x3xf64>) -> tensor<61x3xf64>
    %3999 = stablehlo.broadcast_in_dim %3997, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4000 = stablehlo.broadcast_in_dim %3999, dims = [0, 1] : (tensor<61x1xf64>) -> tensor<61x3xf64>
    %4001 = stablehlo.multiply %4000, %3998 : tensor<61x3xf64>
    %4002 = stablehlo.add %3995, %4001 : tensor<61x3xf64>
    %4003 = stablehlo.add %3972, %4002 : tensor<61x3xf64>
    %4004 = stablehlo.slice %3979 [0:61, 0:1] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4005 = stablehlo.reshape %4004 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4006 = stablehlo.slice %arg8 [0:61, 0:1] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4007 = stablehlo.reshape %4006 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4008 = stablehlo.multiply %4005, %4007 : tensor<61xf64>
    %4009 = stablehlo.slice %3979 [0:61, 1:2] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4010 = stablehlo.reshape %4009 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4011 = stablehlo.slice %arg8 [0:61, 1:2] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4012 = stablehlo.reshape %4011 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4013 = stablehlo.multiply %4010, %4012 : tensor<61xf64>
    %4014 = stablehlo.subtract %4008, %4013 : tensor<61xf64>
    %4015 = stablehlo.slice %3979 [0:61, 2:3] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4016 = stablehlo.reshape %4015 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4017 = stablehlo.slice %arg8 [0:61, 2:3] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4018 = stablehlo.reshape %4017 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4019 = stablehlo.multiply %4016, %4018 : tensor<61xf64>
    %4020 = stablehlo.subtract %4014, %4019 : tensor<61xf64>
    %4021 = stablehlo.slice %3979 [0:61, 3:4] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4022 = stablehlo.reshape %4021 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4023 = stablehlo.slice %arg8 [0:61, 3:4] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4024 = stablehlo.reshape %4023 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4025 = stablehlo.multiply %4022, %4024 : tensor<61xf64>
    %4026 = stablehlo.subtract %4020, %4025 : tensor<61xf64>
    %4027 = stablehlo.slice %3979 [0:61, 0:1] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4028 = stablehlo.reshape %4027 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4029 = stablehlo.slice %arg8 [0:61, 1:2] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4030 = stablehlo.reshape %4029 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4031 = stablehlo.multiply %4028, %4030 : tensor<61xf64>
    %4032 = stablehlo.slice %3979 [0:61, 1:2] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4033 = stablehlo.reshape %4032 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4034 = stablehlo.slice %arg8 [0:61, 0:1] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4035 = stablehlo.reshape %4034 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4036 = stablehlo.multiply %4033, %4035 : tensor<61xf64>
    %4037 = stablehlo.add %4031, %4036 : tensor<61xf64>
    %4038 = stablehlo.slice %3979 [0:61, 2:3] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4039 = stablehlo.reshape %4038 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4040 = stablehlo.slice %arg8 [0:61, 3:4] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4041 = stablehlo.reshape %4040 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4042 = stablehlo.multiply %4039, %4041 : tensor<61xf64>
    %4043 = stablehlo.add %4037, %4042 : tensor<61xf64>
    %4044 = stablehlo.slice %3979 [0:61, 3:4] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4045 = stablehlo.reshape %4044 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4046 = stablehlo.slice %arg8 [0:61, 2:3] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4047 = stablehlo.reshape %4046 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4048 = stablehlo.multiply %4045, %4047 : tensor<61xf64>
    %4049 = stablehlo.subtract %4043, %4048 : tensor<61xf64>
    %4050 = stablehlo.slice %3979 [0:61, 0:1] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4051 = stablehlo.reshape %4050 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4052 = stablehlo.slice %arg8 [0:61, 2:3] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4053 = stablehlo.reshape %4052 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4054 = stablehlo.multiply %4051, %4053 : tensor<61xf64>
    %4055 = stablehlo.slice %3979 [0:61, 1:2] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4056 = stablehlo.reshape %4055 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4057 = stablehlo.slice %arg8 [0:61, 3:4] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4058 = stablehlo.reshape %4057 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4059 = stablehlo.multiply %4056, %4058 : tensor<61xf64>
    %4060 = stablehlo.subtract %4054, %4059 : tensor<61xf64>
    %4061 = stablehlo.slice %3979 [0:61, 2:3] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4062 = stablehlo.reshape %4061 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4063 = stablehlo.slice %arg8 [0:61, 0:1] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4064 = stablehlo.reshape %4063 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4065 = stablehlo.multiply %4062, %4064 : tensor<61xf64>
    %4066 = stablehlo.add %4060, %4065 : tensor<61xf64>
    %4067 = stablehlo.slice %3979 [0:61, 3:4] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4068 = stablehlo.reshape %4067 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4069 = stablehlo.slice %arg8 [0:61, 1:2] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4070 = stablehlo.reshape %4069 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4071 = stablehlo.multiply %4068, %4070 : tensor<61xf64>
    %4072 = stablehlo.add %4066, %4071 : tensor<61xf64>
    %4073 = stablehlo.slice %3979 [0:61, 0:1] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4074 = stablehlo.reshape %4073 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4075 = stablehlo.slice %arg8 [0:61, 3:4] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4076 = stablehlo.reshape %4075 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4077 = stablehlo.multiply %4074, %4076 : tensor<61xf64>
    %4078 = stablehlo.slice %3979 [0:61, 1:2] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4079 = stablehlo.reshape %4078 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4080 = stablehlo.slice %arg8 [0:61, 2:3] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4081 = stablehlo.reshape %4080 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4082 = stablehlo.multiply %4079, %4081 : tensor<61xf64>
    %4083 = stablehlo.add %4077, %4082 : tensor<61xf64>
    %4084 = stablehlo.slice %3979 [0:61, 2:3] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4085 = stablehlo.reshape %4084 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4086 = stablehlo.slice %arg8 [0:61, 1:2] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4087 = stablehlo.reshape %4086 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4088 = stablehlo.multiply %4085, %4087 : tensor<61xf64>
    %4089 = stablehlo.subtract %4083, %4088 : tensor<61xf64>
    %4090 = stablehlo.slice %3979 [0:61, 3:4] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4091 = stablehlo.reshape %4090 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4092 = stablehlo.slice %arg8 [0:61, 0:1] : (tensor<61x4xf64>) -> tensor<61x1xf64>
    %4093 = stablehlo.reshape %4092 : (tensor<61x1xf64>) -> tensor<61xf64>
    %4094 = stablehlo.multiply %4091, %4093 : tensor<61xf64>
    %4095 = stablehlo.add %4089, %4094 : tensor<61xf64>
    %4096 = stablehlo.broadcast_in_dim %4026, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4097 = stablehlo.broadcast_in_dim %4049, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4098 = stablehlo.broadcast_in_dim %4072, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4099 = stablehlo.broadcast_in_dim %4095, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4100 = stablehlo.concatenate %4096, %4097, %4098, %4099, dim = 1 : (tensor<61x1xf64>, tensor<61x1xf64>, tensor<61x1xf64>, tensor<61x1xf64>) -> tensor<61x4xf64>
    %4101 = stablehlo.broadcast_in_dim %4100, dims = [0, 1] : (tensor<61x4xf64>) -> tensor<61x4x1xf64>
    %4102 = stablehlo.broadcast_in_dim %4100, dims = [0, 2] : (tensor<61x4xf64>) -> tensor<61x1x4xf64>
    %4103 = stablehlo.broadcast_in_dim %4101, dims = [0, 1, 2] : (tensor<61x4x1xf64>) -> tensor<61x4x4xf64>
    %4104 = stablehlo.broadcast_in_dim %4102, dims = [0, 1, 2] : (tensor<61x1x4xf64>) -> tensor<61x4x4xf64>
    %4105 = stablehlo.multiply %4103, %4104 : tensor<61x4x4xf64>
    %4106 = stablehlo.slice %4105 [0:61, 0:1, 0:1] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4107 = stablehlo.reshape %4106 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4108 = stablehlo.slice %4105 [0:61, 1:2, 1:2] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4109 = stablehlo.reshape %4108 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4110 = stablehlo.add %4107, %4109 : tensor<61xf64>
    %4111 = stablehlo.slice %4105 [0:61, 2:3, 2:3] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4112 = stablehlo.reshape %4111 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4113 = stablehlo.subtract %4110, %4112 : tensor<61xf64>
    %4114 = stablehlo.slice %4105 [0:61, 3:4, 3:4] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4115 = stablehlo.reshape %4114 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4116 = stablehlo.subtract %4113, %4115 : tensor<61xf64>
    %4117 = stablehlo.slice %4105 [0:61, 1:2, 2:3] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4118 = stablehlo.reshape %4117 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4119 = stablehlo.slice %4105 [0:61, 0:1, 3:4] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4120 = stablehlo.reshape %4119 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4121 = stablehlo.subtract %4118, %4120 : tensor<61xf64>
    %cst_348 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4122 = stablehlo.broadcast_in_dim %cst_348, dims = [] : (tensor<f64>) -> tensor<61xf64>
    %4123 = stablehlo.multiply %4122, %4121 : tensor<61xf64>
    %4124 = stablehlo.slice %4105 [0:61, 1:2, 3:4] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4125 = stablehlo.reshape %4124 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4126 = stablehlo.slice %4105 [0:61, 0:1, 2:3] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4127 = stablehlo.reshape %4126 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4128 = stablehlo.add %4125, %4127 : tensor<61xf64>
    %cst_349 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4129 = stablehlo.broadcast_in_dim %cst_349, dims = [] : (tensor<f64>) -> tensor<61xf64>
    %4130 = stablehlo.multiply %4129, %4128 : tensor<61xf64>
    %4131 = stablehlo.slice %4105 [0:61, 1:2, 2:3] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4132 = stablehlo.reshape %4131 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4133 = stablehlo.slice %4105 [0:61, 0:1, 3:4] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4134 = stablehlo.reshape %4133 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4135 = stablehlo.add %4132, %4134 : tensor<61xf64>
    %cst_350 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4136 = stablehlo.broadcast_in_dim %cst_350, dims = [] : (tensor<f64>) -> tensor<61xf64>
    %4137 = stablehlo.multiply %4136, %4135 : tensor<61xf64>
    %4138 = stablehlo.slice %4105 [0:61, 0:1, 0:1] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4139 = stablehlo.reshape %4138 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4140 = stablehlo.slice %4105 [0:61, 1:2, 1:2] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4141 = stablehlo.reshape %4140 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4142 = stablehlo.subtract %4139, %4141 : tensor<61xf64>
    %4143 = stablehlo.slice %4105 [0:61, 2:3, 2:3] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4144 = stablehlo.reshape %4143 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4145 = stablehlo.add %4142, %4144 : tensor<61xf64>
    %4146 = stablehlo.slice %4105 [0:61, 3:4, 3:4] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4147 = stablehlo.reshape %4146 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4148 = stablehlo.subtract %4145, %4147 : tensor<61xf64>
    %4149 = stablehlo.slice %4105 [0:61, 2:3, 3:4] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4150 = stablehlo.reshape %4149 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4151 = stablehlo.slice %4105 [0:61, 0:1, 1:2] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4152 = stablehlo.reshape %4151 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4153 = stablehlo.subtract %4150, %4152 : tensor<61xf64>
    %cst_351 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4154 = stablehlo.broadcast_in_dim %cst_351, dims = [] : (tensor<f64>) -> tensor<61xf64>
    %4155 = stablehlo.multiply %4154, %4153 : tensor<61xf64>
    %4156 = stablehlo.slice %4105 [0:61, 1:2, 3:4] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4157 = stablehlo.reshape %4156 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4158 = stablehlo.slice %4105 [0:61, 0:1, 2:3] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4159 = stablehlo.reshape %4158 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4160 = stablehlo.subtract %4157, %4159 : tensor<61xf64>
    %cst_352 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4161 = stablehlo.broadcast_in_dim %cst_352, dims = [] : (tensor<f64>) -> tensor<61xf64>
    %4162 = stablehlo.multiply %4161, %4160 : tensor<61xf64>
    %4163 = stablehlo.slice %4105 [0:61, 2:3, 3:4] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4164 = stablehlo.reshape %4163 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4165 = stablehlo.slice %4105 [0:61, 0:1, 1:2] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4166 = stablehlo.reshape %4165 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4167 = stablehlo.add %4164, %4166 : tensor<61xf64>
    %cst_353 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4168 = stablehlo.broadcast_in_dim %cst_353, dims = [] : (tensor<f64>) -> tensor<61xf64>
    %4169 = stablehlo.multiply %4168, %4167 : tensor<61xf64>
    %4170 = stablehlo.slice %4105 [0:61, 0:1, 0:1] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4171 = stablehlo.reshape %4170 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4172 = stablehlo.slice %4105 [0:61, 1:2, 1:2] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4173 = stablehlo.reshape %4172 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4174 = stablehlo.subtract %4171, %4173 : tensor<61xf64>
    %4175 = stablehlo.slice %4105 [0:61, 2:3, 2:3] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4176 = stablehlo.reshape %4175 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4177 = stablehlo.subtract %4174, %4176 : tensor<61xf64>
    %4178 = stablehlo.slice %4105 [0:61, 3:4, 3:4] : (tensor<61x4x4xf64>) -> tensor<61x1x1xf64>
    %4179 = stablehlo.reshape %4178 : (tensor<61x1x1xf64>) -> tensor<61xf64>
    %4180 = stablehlo.add %4177, %4179 : tensor<61xf64>
    %4181 = stablehlo.broadcast_in_dim %4116, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4182 = stablehlo.broadcast_in_dim %4123, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4183 = stablehlo.broadcast_in_dim %4130, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4184 = stablehlo.concatenate %4181, %4182, %4183, dim = 1 : (tensor<61x1xf64>, tensor<61x1xf64>, tensor<61x1xf64>) -> tensor<61x3xf64>
    %4185 = stablehlo.broadcast_in_dim %4137, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4186 = stablehlo.broadcast_in_dim %4148, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4187 = stablehlo.broadcast_in_dim %4155, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4188 = stablehlo.concatenate %4185, %4186, %4187, dim = 1 : (tensor<61x1xf64>, tensor<61x1xf64>, tensor<61x1xf64>) -> tensor<61x3xf64>
    %4189 = stablehlo.broadcast_in_dim %4162, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4190 = stablehlo.broadcast_in_dim %4169, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4191 = stablehlo.broadcast_in_dim %4180, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %4192 = stablehlo.concatenate %4189, %4190, %4191, dim = 1 : (tensor<61x1xf64>, tensor<61x1xf64>, tensor<61x1xf64>) -> tensor<61x3xf64>
    %4193 = stablehlo.broadcast_in_dim %4184, dims = [0, 2] : (tensor<61x3xf64>) -> tensor<61x1x3xf64>
    %4194 = stablehlo.broadcast_in_dim %4188, dims = [0, 2] : (tensor<61x3xf64>) -> tensor<61x1x3xf64>
    %4195 = stablehlo.broadcast_in_dim %4192, dims = [0, 2] : (tensor<61x3xf64>) -> tensor<61x1x3xf64>
    %4196 = stablehlo.concatenate %4193, %4194, %4195, dim = 1 : (tensor<61x1x3xf64>, tensor<61x1x3xf64>, tensor<61x1x3xf64>) -> tensor<61x3x3xf64>
    %c_354 = stablehlo.constant dense<0> : tensor<i32>
    %4197 = stablehlo.broadcast_in_dim %c_354, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %4198 = stablehlo.compare  LT, %c_43, %4197,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %c_355 = stablehlo.constant dense<9> : tensor<i32>
    %4199 = stablehlo.broadcast_in_dim %c_355, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %4200 = stablehlo.add %c_43, %4199 : tensor<1xi32>
    %4201 = stablehlo.select %4198, %4200, %c_43 : tensor<1xi1>, tensor<1xi32>
    %4202 = stablehlo.broadcast_in_dim %4201, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %4203 = "stablehlo.gather"(%3744, %4202) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 3>}> : (tensor<9x3xf64>, tensor<1x1xi32>) -> tensor<1x3xf64>
    %c_356 = stablehlo.constant dense<0> : tensor<i32>
    %4204 = stablehlo.broadcast_in_dim %c_356, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %4205 = stablehlo.compare  LT, %c_43, %4204,  SIGNED : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %c_357 = stablehlo.constant dense<9> : tensor<i32>
    %4206 = stablehlo.broadcast_in_dim %c_357, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %4207 = stablehlo.add %c_43, %4206 : tensor<1xi32>
    %4208 = stablehlo.select %4205, %4207, %c_43 : tensor<1xi1>, tensor<1xi32>
    %4209 = stablehlo.broadcast_in_dim %4208, dims = [0] : (tensor<1xi32>) -> tensor<1x1xi32>
    %4210 = "stablehlo.gather"(%3746, %4209) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 4>}> : (tensor<9x4xf64>, tensor<1x1xi32>) -> tensor<1x4xf64>
    %4211 = stablehlo.slice %4210 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4212 = stablehlo.reshape %4211 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4213 = stablehlo.slice %4210 [0:1, 1:4] : (tensor<1x4xf64>) -> tensor<1x3xf64>
    %4214 = stablehlo.dot_general %4213, %arg9, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %4215 = stablehlo.broadcast_in_dim %4214, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4216 = stablehlo.broadcast_in_dim %4215, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %4217 = stablehlo.multiply %4216, %4213 : tensor<1x3xf64>
    %cst_358 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4218 = stablehlo.broadcast_in_dim %cst_358, dims = [] : (tensor<f64>) -> tensor<1x3xf64>
    %4219 = stablehlo.multiply %4218, %4217 : tensor<1x3xf64>
    %4220 = stablehlo.multiply %4212, %4212 : tensor<1xf64>
    %4221 = stablehlo.dot_general %4213, %4213, batching_dims = [0] x [0], contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1xf64>
    %4222 = stablehlo.subtract %4220, %4221 : tensor<1xf64>
    %4223 = stablehlo.broadcast_in_dim %4222, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4224 = stablehlo.broadcast_in_dim %4223, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %4225 = stablehlo.multiply %4224, %arg9 : tensor<1x3xf64>
    %4226 = stablehlo.add %4219, %4225 : tensor<1x3xf64>
    %cst_359 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4227 = stablehlo.broadcast_in_dim %cst_359, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %4228 = stablehlo.multiply %4227, %4212 : tensor<1xf64>
    %4229 = call @cross(%4213, %arg9) : (tensor<1x3xf64>, tensor<1x3xf64>) -> tensor<1x3xf64>
    %4230 = stablehlo.broadcast_in_dim %4228, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4231 = stablehlo.broadcast_in_dim %4230, dims = [0, 1] : (tensor<1x1xf64>) -> tensor<1x3xf64>
    %4232 = stablehlo.multiply %4231, %4229 : tensor<1x3xf64>
    %4233 = stablehlo.add %4226, %4232 : tensor<1x3xf64>
    %4234 = stablehlo.add %4203, %4233 : tensor<1x3xf64>
    %4235 = stablehlo.slice %4210 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4236 = stablehlo.reshape %4235 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4237 = stablehlo.slice %arg10 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4238 = stablehlo.reshape %4237 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4239 = stablehlo.multiply %4236, %4238 : tensor<1xf64>
    %4240 = stablehlo.slice %4210 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4241 = stablehlo.reshape %4240 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4242 = stablehlo.slice %arg10 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4243 = stablehlo.reshape %4242 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4244 = stablehlo.multiply %4241, %4243 : tensor<1xf64>
    %4245 = stablehlo.subtract %4239, %4244 : tensor<1xf64>
    %4246 = stablehlo.slice %4210 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4247 = stablehlo.reshape %4246 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4248 = stablehlo.slice %arg10 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4249 = stablehlo.reshape %4248 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4250 = stablehlo.multiply %4247, %4249 : tensor<1xf64>
    %4251 = stablehlo.subtract %4245, %4250 : tensor<1xf64>
    %4252 = stablehlo.slice %4210 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4253 = stablehlo.reshape %4252 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4254 = stablehlo.slice %arg10 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4255 = stablehlo.reshape %4254 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4256 = stablehlo.multiply %4253, %4255 : tensor<1xf64>
    %4257 = stablehlo.subtract %4251, %4256 : tensor<1xf64>
    %4258 = stablehlo.slice %4210 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4259 = stablehlo.reshape %4258 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4260 = stablehlo.slice %arg10 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4261 = stablehlo.reshape %4260 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4262 = stablehlo.multiply %4259, %4261 : tensor<1xf64>
    %4263 = stablehlo.slice %4210 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4264 = stablehlo.reshape %4263 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4265 = stablehlo.slice %arg10 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4266 = stablehlo.reshape %4265 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4267 = stablehlo.multiply %4264, %4266 : tensor<1xf64>
    %4268 = stablehlo.add %4262, %4267 : tensor<1xf64>
    %4269 = stablehlo.slice %4210 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4270 = stablehlo.reshape %4269 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4271 = stablehlo.slice %arg10 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4272 = stablehlo.reshape %4271 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4273 = stablehlo.multiply %4270, %4272 : tensor<1xf64>
    %4274 = stablehlo.add %4268, %4273 : tensor<1xf64>
    %4275 = stablehlo.slice %4210 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4276 = stablehlo.reshape %4275 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4277 = stablehlo.slice %arg10 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4278 = stablehlo.reshape %4277 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4279 = stablehlo.multiply %4276, %4278 : tensor<1xf64>
    %4280 = stablehlo.subtract %4274, %4279 : tensor<1xf64>
    %4281 = stablehlo.slice %4210 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4282 = stablehlo.reshape %4281 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4283 = stablehlo.slice %arg10 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4284 = stablehlo.reshape %4283 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4285 = stablehlo.multiply %4282, %4284 : tensor<1xf64>
    %4286 = stablehlo.slice %4210 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4287 = stablehlo.reshape %4286 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4288 = stablehlo.slice %arg10 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4289 = stablehlo.reshape %4288 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4290 = stablehlo.multiply %4287, %4289 : tensor<1xf64>
    %4291 = stablehlo.subtract %4285, %4290 : tensor<1xf64>
    %4292 = stablehlo.slice %4210 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4293 = stablehlo.reshape %4292 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4294 = stablehlo.slice %arg10 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4295 = stablehlo.reshape %4294 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4296 = stablehlo.multiply %4293, %4295 : tensor<1xf64>
    %4297 = stablehlo.add %4291, %4296 : tensor<1xf64>
    %4298 = stablehlo.slice %4210 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4299 = stablehlo.reshape %4298 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4300 = stablehlo.slice %arg10 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4301 = stablehlo.reshape %4300 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4302 = stablehlo.multiply %4299, %4301 : tensor<1xf64>
    %4303 = stablehlo.add %4297, %4302 : tensor<1xf64>
    %4304 = stablehlo.slice %4210 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4305 = stablehlo.reshape %4304 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4306 = stablehlo.slice %arg10 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4307 = stablehlo.reshape %4306 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4308 = stablehlo.multiply %4305, %4307 : tensor<1xf64>
    %4309 = stablehlo.slice %4210 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4310 = stablehlo.reshape %4309 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4311 = stablehlo.slice %arg10 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4312 = stablehlo.reshape %4311 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4313 = stablehlo.multiply %4310, %4312 : tensor<1xf64>
    %4314 = stablehlo.add %4308, %4313 : tensor<1xf64>
    %4315 = stablehlo.slice %4210 [0:1, 2:3] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4316 = stablehlo.reshape %4315 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4317 = stablehlo.slice %arg10 [0:1, 1:2] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4318 = stablehlo.reshape %4317 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4319 = stablehlo.multiply %4316, %4318 : tensor<1xf64>
    %4320 = stablehlo.subtract %4314, %4319 : tensor<1xf64>
    %4321 = stablehlo.slice %4210 [0:1, 3:4] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4322 = stablehlo.reshape %4321 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4323 = stablehlo.slice %arg10 [0:1, 0:1] : (tensor<1x4xf64>) -> tensor<1x1xf64>
    %4324 = stablehlo.reshape %4323 : (tensor<1x1xf64>) -> tensor<1xf64>
    %4325 = stablehlo.multiply %4322, %4324 : tensor<1xf64>
    %4326 = stablehlo.add %4320, %4325 : tensor<1xf64>
    %4327 = stablehlo.broadcast_in_dim %4257, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4328 = stablehlo.broadcast_in_dim %4280, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4329 = stablehlo.broadcast_in_dim %4303, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4330 = stablehlo.broadcast_in_dim %4326, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4331 = stablehlo.concatenate %4327, %4328, %4329, %4330, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x4xf64>
    %4332 = stablehlo.broadcast_in_dim %4331, dims = [0, 1] : (tensor<1x4xf64>) -> tensor<1x4x1xf64>
    %4333 = stablehlo.broadcast_in_dim %4331, dims = [0, 2] : (tensor<1x4xf64>) -> tensor<1x1x4xf64>
    %4334 = stablehlo.broadcast_in_dim %4332, dims = [0, 1, 2] : (tensor<1x4x1xf64>) -> tensor<1x4x4xf64>
    %4335 = stablehlo.broadcast_in_dim %4333, dims = [0, 1, 2] : (tensor<1x1x4xf64>) -> tensor<1x4x4xf64>
    %4336 = stablehlo.multiply %4334, %4335 : tensor<1x4x4xf64>
    %4337 = stablehlo.slice %4336 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4338 = stablehlo.reshape %4337 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4339 = stablehlo.slice %4336 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4340 = stablehlo.reshape %4339 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4341 = stablehlo.add %4338, %4340 : tensor<1xf64>
    %4342 = stablehlo.slice %4336 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4343 = stablehlo.reshape %4342 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4344 = stablehlo.subtract %4341, %4343 : tensor<1xf64>
    %4345 = stablehlo.slice %4336 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4346 = stablehlo.reshape %4345 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4347 = stablehlo.subtract %4344, %4346 : tensor<1xf64>
    %4348 = stablehlo.slice %4336 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4349 = stablehlo.reshape %4348 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4350 = stablehlo.slice %4336 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4351 = stablehlo.reshape %4350 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4352 = stablehlo.subtract %4349, %4351 : tensor<1xf64>
    %cst_360 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4353 = stablehlo.broadcast_in_dim %cst_360, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %4354 = stablehlo.multiply %4353, %4352 : tensor<1xf64>
    %4355 = stablehlo.slice %4336 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4356 = stablehlo.reshape %4355 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4357 = stablehlo.slice %4336 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4358 = stablehlo.reshape %4357 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4359 = stablehlo.add %4356, %4358 : tensor<1xf64>
    %cst_361 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4360 = stablehlo.broadcast_in_dim %cst_361, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %4361 = stablehlo.multiply %4360, %4359 : tensor<1xf64>
    %4362 = stablehlo.slice %4336 [0:1, 1:2, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4363 = stablehlo.reshape %4362 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4364 = stablehlo.slice %4336 [0:1, 0:1, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4365 = stablehlo.reshape %4364 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4366 = stablehlo.add %4363, %4365 : tensor<1xf64>
    %cst_362 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4367 = stablehlo.broadcast_in_dim %cst_362, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %4368 = stablehlo.multiply %4367, %4366 : tensor<1xf64>
    %4369 = stablehlo.slice %4336 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4370 = stablehlo.reshape %4369 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4371 = stablehlo.slice %4336 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4372 = stablehlo.reshape %4371 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4373 = stablehlo.subtract %4370, %4372 : tensor<1xf64>
    %4374 = stablehlo.slice %4336 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4375 = stablehlo.reshape %4374 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4376 = stablehlo.add %4373, %4375 : tensor<1xf64>
    %4377 = stablehlo.slice %4336 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4378 = stablehlo.reshape %4377 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4379 = stablehlo.subtract %4376, %4378 : tensor<1xf64>
    %4380 = stablehlo.slice %4336 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4381 = stablehlo.reshape %4380 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4382 = stablehlo.slice %4336 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4383 = stablehlo.reshape %4382 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4384 = stablehlo.subtract %4381, %4383 : tensor<1xf64>
    %cst_363 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4385 = stablehlo.broadcast_in_dim %cst_363, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %4386 = stablehlo.multiply %4385, %4384 : tensor<1xf64>
    %4387 = stablehlo.slice %4336 [0:1, 1:2, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4388 = stablehlo.reshape %4387 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4389 = stablehlo.slice %4336 [0:1, 0:1, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4390 = stablehlo.reshape %4389 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4391 = stablehlo.subtract %4388, %4390 : tensor<1xf64>
    %cst_364 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4392 = stablehlo.broadcast_in_dim %cst_364, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %4393 = stablehlo.multiply %4392, %4391 : tensor<1xf64>
    %4394 = stablehlo.slice %4336 [0:1, 2:3, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4395 = stablehlo.reshape %4394 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4396 = stablehlo.slice %4336 [0:1, 0:1, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4397 = stablehlo.reshape %4396 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4398 = stablehlo.add %4395, %4397 : tensor<1xf64>
    %cst_365 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %4399 = stablehlo.broadcast_in_dim %cst_365, dims = [] : (tensor<f64>) -> tensor<1xf64>
    %4400 = stablehlo.multiply %4399, %4398 : tensor<1xf64>
    %4401 = stablehlo.slice %4336 [0:1, 0:1, 0:1] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4402 = stablehlo.reshape %4401 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4403 = stablehlo.slice %4336 [0:1, 1:2, 1:2] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4404 = stablehlo.reshape %4403 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4405 = stablehlo.subtract %4402, %4404 : tensor<1xf64>
    %4406 = stablehlo.slice %4336 [0:1, 2:3, 2:3] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4407 = stablehlo.reshape %4406 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4408 = stablehlo.subtract %4405, %4407 : tensor<1xf64>
    %4409 = stablehlo.slice %4336 [0:1, 3:4, 3:4] : (tensor<1x4x4xf64>) -> tensor<1x1x1xf64>
    %4410 = stablehlo.reshape %4409 : (tensor<1x1x1xf64>) -> tensor<1xf64>
    %4411 = stablehlo.add %4408, %4410 : tensor<1xf64>
    %4412 = stablehlo.broadcast_in_dim %4347, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4413 = stablehlo.broadcast_in_dim %4354, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4414 = stablehlo.broadcast_in_dim %4361, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4415 = stablehlo.concatenate %4412, %4413, %4414, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %4416 = stablehlo.broadcast_in_dim %4368, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4417 = stablehlo.broadcast_in_dim %4379, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4418 = stablehlo.broadcast_in_dim %4386, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4419 = stablehlo.concatenate %4416, %4417, %4418, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %4420 = stablehlo.broadcast_in_dim %4393, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4421 = stablehlo.broadcast_in_dim %4400, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4422 = stablehlo.broadcast_in_dim %4411, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %4423 = stablehlo.concatenate %4420, %4421, %4422, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    %4424 = stablehlo.broadcast_in_dim %4415, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %4425 = stablehlo.broadcast_in_dim %4419, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %4426 = stablehlo.broadcast_in_dim %4423, dims = [0, 2] : (tensor<1x3xf64>) -> tensor<1x1x3xf64>
    %4427 = stablehlo.concatenate %4424, %4425, %4426, dim = 1 : (tensor<1x1x3xf64>, tensor<1x1x3xf64>, tensor<1x1x3xf64>) -> tensor<1x3x3xf64>
    return %arg11, %arg12, %3720, %arg14, %arg15, %arg16, %arg17, %arg18, %arg19, %arg20, %arg21, %arg22, %arg23, %arg24, %arg25, %arg26, %3744, %3746, %3748, %3772, %3965, %3731, %3742, %4003, %4196, %4234, %4427, %arg27, %arg28, %arg29, %arg30, %arg31, %arg32, %arg33, %arg34, %arg35, %arg36, %arg37, %arg38, %arg39, %arg40, %arg41, %arg42, %arg43, %arg44, %arg45, %arg46, %arg47, %arg48, %arg49, %arg50, %arg51, %arg52, %arg53, %arg54, %arg55, %arg56, %arg57, %arg58, %arg59, %arg60, %arg61, %arg62, %arg63, %arg64, %arg65, %arg66, %arg67, %arg68, %arg69, %arg70, %arg71, %arg72, %arg73, %arg74, %arg75, %arg76, %arg77, %arg78, %arg79, %arg80, %arg81, %arg82, %arg83, %arg84, %arg85, %arg86, %arg87, %arg88, %arg89, %arg90, %arg91, %arg92, %arg93, %arg94, %arg95, %arg96, %arg97, %arg98, %arg99, %arg100, %arg101, %arg102, %arg103, %arg104, %arg105, %arg106, %arg107, %arg108, %arg109, %arg110, %arg111, %arg112 : tensor<i64>, tensor<f64>, tensor<7xf64>, tensor<7xf64>, tensor<0xf64>, tensor<7xf64>, tensor<7xf64>, tensor<7xf64>, tensor<9x6xf64>, tensor<0xui8>, tensor<0xf64>, tensor<0xf64>, tensor<7xf64>, tensor<0xf64>, tensor<0xf64>, tensor<0xf64>, tensor<9x3xf64>, tensor<9x4xf64>, tensor<9x3x3xf64>, tensor<9x3xf64>, tensor<9x3x3xf64>, tensor<7x3xf64>, tensor<7x3xf64>, tensor<61x3xf64>, tensor<61x3x3xf64>, tensor<1x3xf64>, tensor<1x3x3xf64>, tensor<0x3xf64>, tensor<0x3x3xf64>, tensor<0xf64>, tensor<0xf64>, tensor<9x3xf64>, tensor<7x6xf64>, tensor<9x10xf64>, tensor<0xf64>, tensor<0x6xf64>, tensor<0xi32>, tensor<0xi32>, tensor<0xi32>, tensor<0xf64>, tensor<0xf64>, tensor<0xi32>, tensor<0xi32>, tensor<0xi32>, tensor<0xi32>, tensor<0xi32>, tensor<0x7xf64>, tensor<0xf64>, tensor<0xi32>, tensor<0xf64>, tensor<7xf64>, tensor<7x7xf64>, tensor<9x10xf64>, tensor<7x7xf64>, tensor<7x7xf64>, tensor<0xf64>, tensor<7xf64>, tensor<0xf64>, tensor<0xui8>, tensor<0xf64>, tensor<0xf64>, tensor<7xf64>, tensor<9x6xf64>, tensor<7x6xf64>, tensor<7xf64>, tensor<0xf64>, tensor<0xf64>, tensor<7xf64>, tensor<7xf64>, tensor<7xf64>, tensor<9x3xf64>, tensor<9x3xf64>, tensor<0xf64>, tensor<0xf64>, tensor<0xi32>, tensor<0xi32>, tensor<0xi32>, tensor<0xi32>, tensor<0xi32>, tensor<0xi32>, tensor<0xf64>, tensor<0xf64>, tensor<7xf64>, tensor<7xf64>, tensor<7xf64>, tensor<7xf64>, tensor<7xf64>, tensor<7xf64>, tensor<9x6xf64>, tensor<9x6xf64>, tensor<9x6xf64>, tensor<396xf64>, tensor<396x3xf64>, tensor<396x3x3xf64>, tensor<396xf64>, tensor<396x5xf64>, tensor<396x2xf64>, tensor<396x2xf64>, tensor<396x5xf64>, tensor<396xi32>, tensor<396xi32>, tensor<396x2xi32>, tensor<1591xi64>, tensor<1591x7xf64>, tensor<1591xf64>, tensor<1591xf64>, tensor<1591xf64>, tensor<1591xf64>, tensor<1591xf64>, tensor<1591xf64>, tensor<0xf64>, tensor<0xf64>, tensor<0xf64>
  }
  func.func private @_take() -> (tensor<1x0xf64> {mhlo.layout_mode = "default"}) {
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f64>) -> tensor<1x0xf64>
    return %0 : tensor<1x0xf64>
  }
  func.func private @_take_0(%arg0: tensor<9x3xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<1xi64> {mhlo.layout_mode = "default"}) -> (tensor<1x3xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<9> : tensor<i64>
    %0 = call @remainder(%arg1, %c) : (tensor<1xi64>, tensor<i64>) -> tensor<1xi64>
    %1 = stablehlo.broadcast_in_dim %0, dims = [0] : (tensor<1xi64>) -> tensor<1x1xi64>
    %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 3>}> : (tensor<9x3xf64>, tensor<1x1xi64>) -> tensor<1x3xf64>
    return %2 : tensor<1x3xf64>
  }
  func.func private @remainder(%arg0: tensor<1xi64> {mhlo.layout_mode = "default"}, %arg1: tensor<i64> {mhlo.layout_mode = "default"}) -> (tensor<1xi64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<0> : tensor<i64>
    %0 = stablehlo.compare  EQ, %arg1, %c,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
    %c_0 = stablehlo.constant dense<1> : tensor<i64>
    %1 = call @_where(%0, %c_0, %arg1) : (tensor<i1>, tensor<i64>, tensor<i64>) -> tensor<i64>
    %2 = stablehlo.broadcast_in_dim %1, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %3 = stablehlo.remainder %arg0, %2 : tensor<1xi64>
    %c_1 = stablehlo.constant dense<0> : tensor<i64>
    %4 = stablehlo.broadcast_in_dim %c_1, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %5 = stablehlo.compare  NE, %3, %4,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_2 = stablehlo.constant dense<0> : tensor<i64>
    %6 = stablehlo.broadcast_in_dim %c_2, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %7 = stablehlo.compare  LT, %3, %6,  SIGNED : (tensor<1xi64>, tensor<1xi64>) -> tensor<1xi1>
    %c_3 = stablehlo.constant dense<0> : tensor<i64>
    %8 = stablehlo.compare  LT, %1, %c_3,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<i1>) -> tensor<1xi1>
    %10 = stablehlo.compare  NE, %7, %9,  UNSIGNED : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1>
    %11 = stablehlo.and %10, %5 : tensor<1xi1>
    %12 = stablehlo.broadcast_in_dim %1, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %13 = stablehlo.add %3, %12 : tensor<1xi64>
    %14 = stablehlo.select %11, %13, %3 : tensor<1xi1>, tensor<1xi64>
    return %14 : tensor<1xi64>
  }
  func.func private @_where(%arg0: tensor<i1> {mhlo.layout_mode = "default"}, %arg1: tensor<i64> {mhlo.layout_mode = "default"}, %arg2: tensor<i64> {mhlo.layout_mode = "default"}) -> (tensor<i64> {mhlo.layout_mode = "default"}) {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i64>
    return %0 : tensor<i64>
  }
  func.func private @_take_1(%arg0: tensor<9x4xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<1xi64> {mhlo.layout_mode = "default"}) -> (tensor<1x4xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<9> : tensor<i64>
    %0 = call @remainder(%arg1, %c) : (tensor<1xi64>, tensor<i64>) -> tensor<1xi64>
    %1 = stablehlo.broadcast_in_dim %0, dims = [0] : (tensor<1xi64>) -> tensor<1x1xi64>
    %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 4>}> : (tensor<9x4xf64>, tensor<1x1xi64>) -> tensor<1x4xf64>
    return %2 : tensor<1x4xf64>
  }
  func.func private @_take_2(%arg0: tensor<1x3xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<1xi64> {mhlo.layout_mode = "default"}) -> (tensor<1x3xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<1> : tensor<i64>
    %0 = call @remainder(%arg1, %c) : (tensor<1xi64>, tensor<i64>) -> tensor<1xi64>
    %1 = stablehlo.broadcast_in_dim %0, dims = [0] : (tensor<1xi64>) -> tensor<1x1xi64>
    %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 3>}> : (tensor<1x3xf64>, tensor<1x1xi64>) -> tensor<1x3xf64>
    return %2 : tensor<1x3xf64>
  }
  func.func private @_take_3(%arg0: tensor<1x4xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<1xi64> {mhlo.layout_mode = "default"}) -> (tensor<1x4xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<1> : tensor<i64>
    %0 = call @remainder(%arg1, %c) : (tensor<1xi64>, tensor<i64>) -> tensor<1xi64>
    %1 = stablehlo.broadcast_in_dim %0, dims = [0] : (tensor<1xi64>) -> tensor<1x1xi64>
    %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 4>}> : (tensor<1x4xf64>, tensor<1x1xi64>) -> tensor<1x4xf64>
    return %2 : tensor<1x4xf64>
  }
  func.func private @cross(%arg0: tensor<1x3xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<1x3xf64> {mhlo.layout_mode = "default"}) -> (tensor<1x3xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<0> : tensor<i32>
    %0 = stablehlo.broadcast_in_dim %c, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %1 = "stablehlo.gather"(%arg0, %0) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 1, 1>}> : (tensor<1x3xf64>, tensor<1xi32>) -> tensor<1xf64>
    %c_0 = stablehlo.constant dense<1> : tensor<i32>
    %2 = stablehlo.broadcast_in_dim %c_0, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %3 = "stablehlo.gather"(%arg0, %2) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 1, 1>}> : (tensor<1x3xf64>, tensor<1xi32>) -> tensor<1xf64>
    %c_1 = stablehlo.constant dense<2> : tensor<i32>
    %4 = stablehlo.broadcast_in_dim %c_1, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %5 = "stablehlo.gather"(%arg0, %4) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 1, 1>}> : (tensor<1x3xf64>, tensor<1xi32>) -> tensor<1xf64>
    %c_2 = stablehlo.constant dense<0> : tensor<i32>
    %6 = stablehlo.broadcast_in_dim %c_2, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %7 = "stablehlo.gather"(%arg1, %6) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 1, 1>}> : (tensor<1x3xf64>, tensor<1xi32>) -> tensor<1xf64>
    %c_3 = stablehlo.constant dense<1> : tensor<i32>
    %8 = stablehlo.broadcast_in_dim %c_3, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %9 = "stablehlo.gather"(%arg1, %8) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 1, 1>}> : (tensor<1x3xf64>, tensor<1xi32>) -> tensor<1xf64>
    %c_4 = stablehlo.constant dense<2> : tensor<i32>
    %10 = stablehlo.broadcast_in_dim %c_4, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %11 = "stablehlo.gather"(%arg1, %10) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 1, 1>}> : (tensor<1x3xf64>, tensor<1xi32>) -> tensor<1xf64>
    %12 = stablehlo.multiply %3, %11 : tensor<1xf64>
    %13 = stablehlo.multiply %5, %9 : tensor<1xf64>
    %14 = stablehlo.subtract %12, %13 : tensor<1xf64>
    %15 = stablehlo.multiply %5, %7 : tensor<1xf64>
    %16 = stablehlo.multiply %1, %11 : tensor<1xf64>
    %17 = stablehlo.subtract %15, %16 : tensor<1xf64>
    %18 = stablehlo.multiply %1, %9 : tensor<1xf64>
    %19 = stablehlo.multiply %3, %7 : tensor<1xf64>
    %20 = stablehlo.subtract %18, %19 : tensor<1xf64>
    %21 = stablehlo.broadcast_in_dim %14, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %22 = stablehlo.broadcast_in_dim %17, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %23 = stablehlo.broadcast_in_dim %20, dims = [0] : (tensor<1xf64>) -> tensor<1x1xf64>
    %24 = stablehlo.concatenate %21, %22, %23, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x3xf64>
    return %24 : tensor<1x3xf64>
  }
  func.func private @_take_4(%arg0: tensor<7x3xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<1x1xi64> {mhlo.layout_mode = "default"}) -> (tensor<1x1x3xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<7> : tensor<i64>
    %0 = call @remainder_5(%arg1, %c) : (tensor<1x1xi64>, tensor<i64>) -> tensor<1x1xi64>
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1] : (tensor<1x1xi64>) -> tensor<1x1x1xi64>
    %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<offset_dims = [2], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 2>, indices_are_sorted = false, slice_sizes = array<i64: 1, 3>}> : (tensor<7x3xf64>, tensor<1x1x1xi64>) -> tensor<1x1x3xf64>
    return %2 : tensor<1x1x3xf64>
  }
  func.func private @remainder_5(%arg0: tensor<1x1xi64> {mhlo.layout_mode = "default"}, %arg1: tensor<i64> {mhlo.layout_mode = "default"}) -> (tensor<1x1xi64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<0> : tensor<i64>
    %0 = stablehlo.compare  EQ, %arg1, %c,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
    %c_0 = stablehlo.constant dense<1> : tensor<i64>
    %1 = call @_where(%0, %c_0, %arg1) : (tensor<i1>, tensor<i64>, tensor<i64>) -> tensor<i64>
    %2 = stablehlo.broadcast_in_dim %1, dims = [] : (tensor<i64>) -> tensor<1x1xi64>
    %3 = stablehlo.remainder %arg0, %2 : tensor<1x1xi64>
    %c_1 = stablehlo.constant dense<0> : tensor<i64>
    %4 = stablehlo.broadcast_in_dim %c_1, dims = [] : (tensor<i64>) -> tensor<1x1xi64>
    %5 = stablehlo.compare  NE, %3, %4,  SIGNED : (tensor<1x1xi64>, tensor<1x1xi64>) -> tensor<1x1xi1>
    %c_2 = stablehlo.constant dense<0> : tensor<i64>
    %6 = stablehlo.broadcast_in_dim %c_2, dims = [] : (tensor<i64>) -> tensor<1x1xi64>
    %7 = stablehlo.compare  LT, %3, %6,  SIGNED : (tensor<1x1xi64>, tensor<1x1xi64>) -> tensor<1x1xi1>
    %c_3 = stablehlo.constant dense<0> : tensor<i64>
    %8 = stablehlo.compare  LT, %1, %c_3,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<i1>) -> tensor<1x1xi1>
    %10 = stablehlo.compare  NE, %7, %9,  UNSIGNED : (tensor<1x1xi1>, tensor<1x1xi1>) -> tensor<1x1xi1>
    %11 = stablehlo.and %10, %5 : tensor<1x1xi1>
    %12 = stablehlo.broadcast_in_dim %1, dims = [] : (tensor<i64>) -> tensor<1x1xi64>
    %13 = stablehlo.add %3, %12 : tensor<1x1xi64>
    %14 = stablehlo.select %11, %13, %3 : tensor<1x1xi1>, tensor<1x1xi64>
    return %14 : tensor<1x1xi64>
  }
  func.func private @_take_6(%arg0: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<1x1xi64> {mhlo.layout_mode = "default"}) -> (tensor<1x1xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<7> : tensor<i64>
    %0 = call @remainder_5(%arg1, %c) : (tensor<1x1xi64>, tensor<i64>) -> tensor<1x1xi64>
    %1 = stablehlo.broadcast_in_dim %0, dims = [0, 1] : (tensor<1x1xi64>) -> tensor<1x1x1xi64>
    %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 2>, indices_are_sorted = false, slice_sizes = array<i64: 1>}> : (tensor<7xf64>, tensor<1x1x1xi64>) -> tensor<1x1xf64>
    return %2 : tensor<1x1xf64>
  }
  func.func private @_where_7(%arg0: tensor<1xi1> {mhlo.layout_mode = "default"}, %arg1: tensor<1xi64> {mhlo.layout_mode = "default"}, %arg2: tensor<1xi64> {mhlo.layout_mode = "default"}) -> (tensor<1xi64> {mhlo.layout_mode = "default"}) {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<1xi1>, tensor<1xi64>
    return %0 : tensor<1xi64>
  }
  func.func private @clip(%arg0: tensor<1xi64> {mhlo.layout_mode = "default"}, %arg1: tensor<i64> {mhlo.layout_mode = "default"}, %arg2: tensor<i64> {mhlo.layout_mode = "default"}) -> (tensor<1xi64> {mhlo.layout_mode = "default"}) {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %1 = stablehlo.maximum %0, %arg0 : tensor<1xi64>
    %2 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<i64>) -> tensor<1xi64>
    %3 = stablehlo.minimum %2, %1 : tensor<1xi64>
    return %3 : tensor<1xi64>
  }
  func.func private @argsort(%arg0: tensor<1xi64> {mhlo.layout_mode = "default"}) -> (tensor<1xi64> {mhlo.layout_mode = "default"}) {
    %0 = stablehlo.iota dim = 0 : tensor<1xi64>
    %1:2 = "stablehlo.sort"(%arg0, %0) <{dimension = 0 : i64, is_stable = true}> ({
    ^bb0(%arg1: tensor<i64>, %arg2: tensor<i64>, %arg3: tensor<i64>, %arg4: tensor<i64>):
      %2 = stablehlo.compare  LT, %arg1, %arg2,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
      stablehlo.return %2 : tensor<i1>
    }) : (tensor<1xi64>, tensor<1xi64>) -> (tensor<1xi64>, tensor<1xi64>)
    return %1#1 : tensor<1xi64>
  }
  func.func private @_cumulative_reduction(%arg0: tensor<4xi1> {mhlo.layout_mode = "default"}) -> (tensor<4xi64> {mhlo.layout_mode = "default"}) {
    %0 = stablehlo.convert %arg0 : (tensor<4xi1>) -> tensor<4xi64>
    %1 = call @cumsum(%0) : (tensor<4xi64>) -> tensor<4xi64>
    return %1 : tensor<4xi64>
  }
  func.func private @cumsum(%arg0: tensor<4xi64>) -> tensor<4xi64> {
    %c = stablehlo.constant dense<0> : tensor<i64>
    %0 = stablehlo.broadcast_in_dim %c, dims = [] : (tensor<i64>) -> tensor<i64>
    %1 = "stablehlo.reduce_window"(%arg0, %0) <{base_dilations = array<i64: 1>, padding = dense<[[3, 0]]> : tensor<1x2xi64>, window_dilations = array<i64: 1>, window_dimensions = array<i64: 4>, window_strides = array<i64: 1>}> ({
    ^bb0(%arg1: tensor<i64>, %arg2: tensor<i64>):
      %2 = stablehlo.add %arg1, %arg2 : tensor<i64>
      stablehlo.return %2 : tensor<i64>
    }) : (tensor<4xi64>, tensor<i64>) -> tensor<4xi64>
    return %1 : tensor<4xi64>
  }
  func.func private @clip_8(%arg0: tensor<4xi64> {mhlo.layout_mode = "default"}, %arg1: tensor<i64> {mhlo.layout_mode = "default"}) -> (tensor<4xi64> {mhlo.layout_mode = "default"}) {
    %0 = stablehlo.convert %arg1 : tensor<i64>
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<i64>) -> tensor<4xi64>
    %2 = stablehlo.maximum %1, %arg0 : tensor<4xi64>
    return %2 : tensor<4xi64>
  }
  func.func private @_cumulative_reduction_9(%arg0: tensor<3xi64> {mhlo.layout_mode = "default"}) -> (tensor<3xi64> {mhlo.layout_mode = "default"}) {
    %0 = call @cumsum_10(%arg0) : (tensor<3xi64>) -> tensor<3xi64>
    return %0 : tensor<3xi64>
  }
  func.func private @cumsum_10(%arg0: tensor<3xi64>) -> tensor<3xi64> {
    %c = stablehlo.constant dense<0> : tensor<i64>
    %0 = stablehlo.broadcast_in_dim %c, dims = [] : (tensor<i64>) -> tensor<i64>
    %1 = "stablehlo.reduce_window"(%arg0, %0) <{base_dilations = array<i64: 1>, padding = dense<[[2, 0]]> : tensor<1x2xi64>, window_dilations = array<i64: 1>, window_dimensions = array<i64: 3>, window_strides = array<i64: 1>}> ({
    ^bb0(%arg1: tensor<i64>, %arg2: tensor<i64>):
      %2 = stablehlo.add %arg1, %arg2 : tensor<i64>
      stablehlo.return %2 : tensor<i64>
    }) : (tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
    return %1 : tensor<3xi64>
  }
  func.func private @floor_divide(%arg0: tensor<3xi64> {mhlo.layout_mode = "default"}, %arg1: tensor<i64> {mhlo.layout_mode = "default"}) -> (tensor<3xi64> {mhlo.layout_mode = "default"}) {
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %1 = stablehlo.divide %arg0, %0 : tensor<3xi64>
    %2 = stablehlo.sign %arg0 : tensor<3xi64>
    %3 = stablehlo.sign %arg1 : tensor<i64>
    %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %5 = stablehlo.compare  NE, %2, %4,  SIGNED : (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi1>
    %6 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %7 = stablehlo.remainder %arg0, %6 : tensor<3xi64>
    %c = stablehlo.constant dense<0> : tensor<i64>
    %8 = stablehlo.broadcast_in_dim %c, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %9 = stablehlo.compare  NE, %7, %8,  SIGNED : (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi1>
    %10 = stablehlo.and %5, %9 : tensor<3xi1>
    %c_0 = stablehlo.constant dense<1> : tensor<i64>
    %11 = stablehlo.broadcast_in_dim %c_0, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %12 = stablehlo.subtract %1, %11 : tensor<3xi64>
    %13 = call @_where_11(%10, %12, %1) : (tensor<3xi1>, tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
    return %13 : tensor<3xi64>
  }
  func.func private @_where_11(%arg0: tensor<3xi1> {mhlo.layout_mode = "default"}, %arg1: tensor<3xi64> {mhlo.layout_mode = "default"}, %arg2: tensor<3xi64> {mhlo.layout_mode = "default"}) -> (tensor<3xi64> {mhlo.layout_mode = "default"}) {
    %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<3xi1>, tensor<3xi64>
    return %0 : tensor<3xi64>
  }
  func.func private @remainder_12(%arg0: tensor<3xi64> {mhlo.layout_mode = "default"}, %arg1: tensor<i64> {mhlo.layout_mode = "default"}) -> (tensor<3xi64> {mhlo.layout_mode = "default"}) {
    %0 = stablehlo.convert %arg1 : tensor<i64>
    %c = stablehlo.constant dense<0> : tensor<i64>
    %1 = stablehlo.compare  EQ, %0, %c,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
    %c_0 = stablehlo.constant dense<1> : tensor<i64>
    %2 = call @_where(%1, %c_0, %0) : (tensor<i1>, tensor<i64>, tensor<i64>) -> tensor<i64>
    %3 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %4 = stablehlo.remainder %arg0, %3 : tensor<3xi64>
    %c_1 = stablehlo.constant dense<0> : tensor<i64>
    %5 = stablehlo.broadcast_in_dim %c_1, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %6 = stablehlo.compare  NE, %4, %5,  SIGNED : (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi1>
    %c_2 = stablehlo.constant dense<0> : tensor<i64>
    %7 = stablehlo.broadcast_in_dim %c_2, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %8 = stablehlo.compare  LT, %4, %7,  SIGNED : (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi1>
    %c_3 = stablehlo.constant dense<0> : tensor<i64>
    %9 = stablehlo.compare  LT, %2, %c_3,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
    %10 = stablehlo.broadcast_in_dim %9, dims = [] : (tensor<i1>) -> tensor<3xi1>
    %11 = stablehlo.compare  NE, %8, %10,  UNSIGNED : (tensor<3xi1>, tensor<3xi1>) -> tensor<3xi1>
    %12 = stablehlo.and %11, %6 : tensor<3xi1>
    %13 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<i64>) -> tensor<3xi64>
    %14 = stablehlo.add %4, %13 : tensor<3xi64>
    %15 = stablehlo.select %12, %14, %4 : tensor<3xi1>, tensor<3xi64>
    return %15 : tensor<3xi64>
  }
  func.func private @_take_13(%arg0: tensor<7xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<7xi64> {mhlo.layout_mode = "default"}) -> (tensor<7xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<7> : tensor<i64>
    %0 = call @remainder_14(%arg1, %c) : (tensor<7xi64>, tensor<i64>) -> tensor<7xi64>
    %1 = stablehlo.broadcast_in_dim %0, dims = [0] : (tensor<7xi64>) -> tensor<7x1xi64>
    %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1>}> : (tensor<7xf64>, tensor<7x1xi64>) -> tensor<7xf64>
    return %2 : tensor<7xf64>
  }
  func.func private @remainder_14(%arg0: tensor<7xi64> {mhlo.layout_mode = "default"}, %arg1: tensor<i64> {mhlo.layout_mode = "default"}) -> (tensor<7xi64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<0> : tensor<i64>
    %0 = stablehlo.compare  EQ, %arg1, %c,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
    %c_0 = stablehlo.constant dense<1> : tensor<i64>
    %1 = call @_where(%0, %c_0, %arg1) : (tensor<i1>, tensor<i64>, tensor<i64>) -> tensor<i64>
    %2 = stablehlo.broadcast_in_dim %1, dims = [] : (tensor<i64>) -> tensor<7xi64>
    %3 = stablehlo.remainder %arg0, %2 : tensor<7xi64>
    %c_1 = stablehlo.constant dense<0> : tensor<i64>
    %4 = stablehlo.broadcast_in_dim %c_1, dims = [] : (tensor<i64>) -> tensor<7xi64>
    %5 = stablehlo.compare  NE, %3, %4,  SIGNED : (tensor<7xi64>, tensor<7xi64>) -> tensor<7xi1>
    %c_2 = stablehlo.constant dense<0> : tensor<i64>
    %6 = stablehlo.broadcast_in_dim %c_2, dims = [] : (tensor<i64>) -> tensor<7xi64>
    %7 = stablehlo.compare  LT, %3, %6,  SIGNED : (tensor<7xi64>, tensor<7xi64>) -> tensor<7xi1>
    %c_3 = stablehlo.constant dense<0> : tensor<i64>
    %8 = stablehlo.compare  LT, %1, %c_3,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<i1>) -> tensor<7xi1>
    %10 = stablehlo.compare  NE, %7, %9,  UNSIGNED : (tensor<7xi1>, tensor<7xi1>) -> tensor<7xi1>
    %11 = stablehlo.and %10, %5 : tensor<7xi1>
    %12 = stablehlo.broadcast_in_dim %1, dims = [] : (tensor<i64>) -> tensor<7xi64>
    %13 = stablehlo.add %3, %12 : tensor<7xi64>
    %14 = stablehlo.select %11, %13, %3 : tensor<7xi1>, tensor<7xi64>
    return %14 : tensor<7xi64>
  }
  func.func private @_take_15(%arg0: tensor<7x3xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<7xi64> {mhlo.layout_mode = "default"}) -> (tensor<7x3xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<7> : tensor<i64>
    %0 = call @remainder_14(%arg1, %c) : (tensor<7xi64>, tensor<i64>) -> tensor<7xi64>
    %1 = stablehlo.broadcast_in_dim %0, dims = [0] : (tensor<7xi64>) -> tensor<7x1xi64>
    %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 3>}> : (tensor<7x3xf64>, tensor<7x1xi64>) -> tensor<7x3xf64>
    return %2 : tensor<7x3xf64>
  }
  func.func private @_take_16(%arg0: tensor<9x3xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<9xi64> {mhlo.layout_mode = "default"}) -> (tensor<9x3xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<9> : tensor<i64>
    %0 = call @remainder_17(%arg1, %c) : (tensor<9xi64>, tensor<i64>) -> tensor<9xi64>
    %1 = stablehlo.broadcast_in_dim %0, dims = [0] : (tensor<9xi64>) -> tensor<9x1xi64>
    %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 3>}> : (tensor<9x3xf64>, tensor<9x1xi64>) -> tensor<9x3xf64>
    return %2 : tensor<9x3xf64>
  }
  func.func private @remainder_17(%arg0: tensor<9xi64> {mhlo.layout_mode = "default"}, %arg1: tensor<i64> {mhlo.layout_mode = "default"}) -> (tensor<9xi64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<0> : tensor<i64>
    %0 = stablehlo.compare  EQ, %arg1, %c,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
    %c_0 = stablehlo.constant dense<1> : tensor<i64>
    %1 = call @_where(%0, %c_0, %arg1) : (tensor<i1>, tensor<i64>, tensor<i64>) -> tensor<i64>
    %2 = stablehlo.broadcast_in_dim %1, dims = [] : (tensor<i64>) -> tensor<9xi64>
    %3 = stablehlo.remainder %arg0, %2 : tensor<9xi64>
    %c_1 = stablehlo.constant dense<0> : tensor<i64>
    %4 = stablehlo.broadcast_in_dim %c_1, dims = [] : (tensor<i64>) -> tensor<9xi64>
    %5 = stablehlo.compare  NE, %3, %4,  SIGNED : (tensor<9xi64>, tensor<9xi64>) -> tensor<9xi1>
    %c_2 = stablehlo.constant dense<0> : tensor<i64>
    %6 = stablehlo.broadcast_in_dim %c_2, dims = [] : (tensor<i64>) -> tensor<9xi64>
    %7 = stablehlo.compare  LT, %3, %6,  SIGNED : (tensor<9xi64>, tensor<9xi64>) -> tensor<9xi1>
    %c_3 = stablehlo.constant dense<0> : tensor<i64>
    %8 = stablehlo.compare  LT, %1, %c_3,  SIGNED : (tensor<i64>, tensor<i64>) -> tensor<i1>
    %9 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<i1>) -> tensor<9xi1>
    %10 = stablehlo.compare  NE, %7, %9,  UNSIGNED : (tensor<9xi1>, tensor<9xi1>) -> tensor<9xi1>
    %11 = stablehlo.and %10, %5 : tensor<9xi1>
    %12 = stablehlo.broadcast_in_dim %1, dims = [] : (tensor<i64>) -> tensor<9xi64>
    %13 = stablehlo.add %3, %12 : tensor<9xi64>
    %14 = stablehlo.select %11, %13, %3 : tensor<9xi1>, tensor<9xi64>
    return %14 : tensor<9xi64>
  }
  func.func private @_take_18(%arg0: tensor<9x4xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<9xi64> {mhlo.layout_mode = "default"}) -> (tensor<9x4xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<9> : tensor<i64>
    %0 = call @remainder_17(%arg1, %c) : (tensor<9xi64>, tensor<i64>) -> tensor<9xi64>
    %1 = stablehlo.broadcast_in_dim %0, dims = [0] : (tensor<9xi64>) -> tensor<9x1xi64>
    %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 4>}> : (tensor<9x4xf64>, tensor<9x1xi64>) -> tensor<9x4xf64>
    return %2 : tensor<9x4xf64>
  }
  func.func private @_take_19(%arg0: tensor<9x3x3xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<9xi64> {mhlo.layout_mode = "default"}) -> (tensor<9x3x3xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<9> : tensor<i64>
    %0 = call @remainder_17(%arg1, %c) : (tensor<9xi64>, tensor<i64>) -> tensor<9xi64>
    %1 = stablehlo.broadcast_in_dim %0, dims = [0] : (tensor<9xi64>) -> tensor<9x1xi64>
    %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<offset_dims = [1, 2], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 3, 3>}> : (tensor<9x3x3xf64>, tensor<9x1xi64>) -> tensor<9x3x3xf64>
    return %2 : tensor<9x3x3xf64>
  }
  func.func private @cross_20(%arg0: tensor<9x3xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<9x3xf64> {mhlo.layout_mode = "default"}) -> (tensor<9x3xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<0> : tensor<i32>
    %0 = stablehlo.broadcast_in_dim %c, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %1 = "stablehlo.gather"(%arg0, %0) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 9, 1>}> : (tensor<9x3xf64>, tensor<1xi32>) -> tensor<9xf64>
    %c_0 = stablehlo.constant dense<1> : tensor<i32>
    %2 = stablehlo.broadcast_in_dim %c_0, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %3 = "stablehlo.gather"(%arg0, %2) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 9, 1>}> : (tensor<9x3xf64>, tensor<1xi32>) -> tensor<9xf64>
    %c_1 = stablehlo.constant dense<2> : tensor<i32>
    %4 = stablehlo.broadcast_in_dim %c_1, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %5 = "stablehlo.gather"(%arg0, %4) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 9, 1>}> : (tensor<9x3xf64>, tensor<1xi32>) -> tensor<9xf64>
    %c_2 = stablehlo.constant dense<0> : tensor<i32>
    %6 = stablehlo.broadcast_in_dim %c_2, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %7 = "stablehlo.gather"(%arg1, %6) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 9, 1>}> : (tensor<9x3xf64>, tensor<1xi32>) -> tensor<9xf64>
    %c_3 = stablehlo.constant dense<1> : tensor<i32>
    %8 = stablehlo.broadcast_in_dim %c_3, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %9 = "stablehlo.gather"(%arg1, %8) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 9, 1>}> : (tensor<9x3xf64>, tensor<1xi32>) -> tensor<9xf64>
    %c_4 = stablehlo.constant dense<2> : tensor<i32>
    %10 = stablehlo.broadcast_in_dim %c_4, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %11 = "stablehlo.gather"(%arg1, %10) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 9, 1>}> : (tensor<9x3xf64>, tensor<1xi32>) -> tensor<9xf64>
    %12 = stablehlo.multiply %3, %11 : tensor<9xf64>
    %13 = stablehlo.multiply %5, %9 : tensor<9xf64>
    %14 = stablehlo.subtract %12, %13 : tensor<9xf64>
    %15 = stablehlo.multiply %5, %7 : tensor<9xf64>
    %16 = stablehlo.multiply %1, %11 : tensor<9xf64>
    %17 = stablehlo.subtract %15, %16 : tensor<9xf64>
    %18 = stablehlo.multiply %1, %9 : tensor<9xf64>
    %19 = stablehlo.multiply %3, %7 : tensor<9xf64>
    %20 = stablehlo.subtract %18, %19 : tensor<9xf64>
    %21 = stablehlo.broadcast_in_dim %14, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %22 = stablehlo.broadcast_in_dim %17, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %23 = stablehlo.broadcast_in_dim %20, dims = [0] : (tensor<9xf64>) -> tensor<9x1xf64>
    %24 = stablehlo.concatenate %21, %22, %23, dim = 1 : (tensor<9x1xf64>, tensor<9x1xf64>, tensor<9x1xf64>) -> tensor<9x3xf64>
    return %24 : tensor<9x3xf64>
  }
  func.func private @cross_21(%arg0: tensor<61x3xf64> {mhlo.layout_mode = "default"}, %arg1: tensor<61x3xf64> {mhlo.layout_mode = "default"}) -> (tensor<61x3xf64> {mhlo.layout_mode = "default"}) {
    %c = stablehlo.constant dense<0> : tensor<i32>
    %0 = stablehlo.broadcast_in_dim %c, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %1 = "stablehlo.gather"(%arg0, %0) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 61, 1>}> : (tensor<61x3xf64>, tensor<1xi32>) -> tensor<61xf64>
    %c_0 = stablehlo.constant dense<1> : tensor<i32>
    %2 = stablehlo.broadcast_in_dim %c_0, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %3 = "stablehlo.gather"(%arg0, %2) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 61, 1>}> : (tensor<61x3xf64>, tensor<1xi32>) -> tensor<61xf64>
    %c_1 = stablehlo.constant dense<2> : tensor<i32>
    %4 = stablehlo.broadcast_in_dim %c_1, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %5 = "stablehlo.gather"(%arg0, %4) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 61, 1>}> : (tensor<61x3xf64>, tensor<1xi32>) -> tensor<61xf64>
    %c_2 = stablehlo.constant dense<0> : tensor<i32>
    %6 = stablehlo.broadcast_in_dim %c_2, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %7 = "stablehlo.gather"(%arg1, %6) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 61, 1>}> : (tensor<61x3xf64>, tensor<1xi32>) -> tensor<61xf64>
    %c_3 = stablehlo.constant dense<1> : tensor<i32>
    %8 = stablehlo.broadcast_in_dim %c_3, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %9 = "stablehlo.gather"(%arg1, %8) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 61, 1>}> : (tensor<61x3xf64>, tensor<1xi32>) -> tensor<61xf64>
    %c_4 = stablehlo.constant dense<2> : tensor<i32>
    %10 = stablehlo.broadcast_in_dim %c_4, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %11 = "stablehlo.gather"(%arg1, %10) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true, slice_sizes = array<i64: 61, 1>}> : (tensor<61x3xf64>, tensor<1xi32>) -> tensor<61xf64>
    %12 = stablehlo.multiply %3, %11 : tensor<61xf64>
    %13 = stablehlo.multiply %5, %9 : tensor<61xf64>
    %14 = stablehlo.subtract %12, %13 : tensor<61xf64>
    %15 = stablehlo.multiply %5, %7 : tensor<61xf64>
    %16 = stablehlo.multiply %1, %11 : tensor<61xf64>
    %17 = stablehlo.subtract %15, %16 : tensor<61xf64>
    %18 = stablehlo.multiply %1, %9 : tensor<61xf64>
    %19 = stablehlo.multiply %3, %7 : tensor<61xf64>
    %20 = stablehlo.subtract %18, %19 : tensor<61xf64>
    %21 = stablehlo.broadcast_in_dim %14, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %22 = stablehlo.broadcast_in_dim %17, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %23 = stablehlo.broadcast_in_dim %20, dims = [0] : (tensor<61xf64>) -> tensor<61x1xf64>
    %24 = stablehlo.concatenate %21, %22, %23, dim = 1 : (tensor<61x1xf64>, tensor<61x1xf64>, tensor<61x1xf64>) -> tensor<61x3xf64>
    return %24 : tensor<61x3xf64>
  }
}