jax.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 147,
   "id": "0c89b9ab-9d96-4990-a3bf-a37e3c9aff40",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "import jax.random\n",
    "import jax.lax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 215,
   "id": "bc23c2f4-4d5b-4eb8-ae36-7e4158d77d02",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rmsnorm(x, weight):\n",
    "  ss = 1 / jnp.sqrt(x.dot(x) / x.shape[0] + 1e-5)\n",
    "  return weight * x * ss\n",
    "\n",
    "def softmax(x):\n",
    "  max_val = jnp.max(x)\n",
    "  x = jnp.exp(x - max_val)\n",
    "  return x / sum(x)\n",
    "\n",
    "def sigmoid(x):\n",
    "  return 1 / (1 + jnp.exp(-x))\n",
    "\n",
    "def silu(x):\n",
    "  return x * sigmoid(x)\n",
    "\n",
    "\n",
    "# Token is token value\n",
    "asserts = False\n",
    "def forward(token, config, weights, key_cache, value_cache):\n",
    "  pos = key_cache.shape[1]\n",
    "  assert pos == key_cache.shape[1]\n",
    "  assert pos == value_cache.shape[1]\n",
    "    \n",
    "  n_layers = config['n_layers']\n",
    "  seq_len = config['seq_len']\n",
    "  n_heads = config['n_heads']\n",
    "  vocab_size = config['vocab_size']\n",
    "\n",
    "  # Total number of parameters of the recurrent state\n",
    "  dim = config['dim']\n",
    "\n",
    "  n_kv_heads = config['n_kv_heads']\n",
    "\n",
    "  # number of hidden dimensions?\n",
    "  hidden_dim = config['hidden_dim']\n",
    "\n",
    "\n",
    "  # Number of parameters per head\n",
    "  head_size = dim // n_heads\n",
    "\n",
    "  # Number of heads per kv\n",
    "  kv_mul = n_heads // n_kv_heads\n",
    "\n",
    "  # Number of parameters in a kv\n",
    "  kv_dim = dim // n_heads * n_kv_heads\n",
    "\n",
    "\n",
    "  wo = weights['wo']\n",
    "  if asserts: assert wo.shape == (n_layers, dim, dim)\n",
    "  rms_ffn_weight = weights['rms_ffn_weight']\n",
    "  if asserts: assert rms_ffn_weight.shape == (n_layers, dim)\n",
    "  w1 = weights['w1']\n",
    "  if asserts: assert w1.shape == (n_layers, hidden_dim, dim)\n",
    "  w3 = weights['w3']\n",
    "  if asserts: assert w3.shape == (n_layers, hidden_dim, dim)\n",
    "  w2 = weights['w2']\n",
    "  if asserts: assert w2.shape == (n_layers, dim, hidden_dim)\n",
    "\n",
    "  rms_att_weight = weights['rms_att_weight']\n",
    "  if asserts: assert rms_att_weight.shape == (n_layers,dim)\n",
    "\n",
    "  rms_final_weight = weights['rms_final_weight']\n",
    "  if asserts: assert rms_final_weight.shape == (dim,)\n",
    "  wcls = weights['wcls']\n",
    "  if asserts: assert wcls.shape == (vocab_size, dim)\n",
    "\n",
    "  token_embedding_table = weights['token_embedding_table']\n",
    "  if asserts: assert token_embedding_table.shape == (vocab_size, dim)\n",
    "\n",
    "  x = token_embedding_table[token, :]\n",
    "  if asserts: assert x.shape == (dim, )\n",
    "\n",
    "  wq = weights['wq']\n",
    "  if asserts: assert wq.shape == (n_layers, dim, dim)\n",
    "\n",
    "  wk = weights['wk']\n",
    "  if asserts: assert wk.shape == (n_layers, kv_dim, dim)\n",
    "\n",
    "  wv = weights['wv']\n",
    "  if asserts: assert wv.shape == (n_layers, kv_dim, dim)\n",
    "\n",
    "  toconv = []\n",
    "       \n",
    "  for i in range(0, dim, 2):\n",
    "    freq = 1 / jnp.power(10000, (i % head_size) / head_size)\n",
    "    val = pos * freq\n",
    "    fcr = jnp.cos(val)\n",
    "    fci = jnp.sin(val)\n",
    "\n",
    "    rotM = jnp.array([[fcr, -fci],\n",
    "                      [fci, fcr]])\n",
    "    toconv.append(rotM)\n",
    "  toconv2 = toconv[:kv_dim//2] + [jnp.eye(2)] * (dim//2 - kv_dim//2)\n",
    "    \n",
    "  toconv = jnp.array(toconv)\n",
    "  toconv2 = jnp.array(toconv2)\n",
    "\n",
    "  keys2 = []\n",
    "  values2 = []\n",
    "  for l in range(n_layers):\n",
    "    xb = rmsnorm(x, rms_att_weight[l, :])\n",
    "    if asserts: assert xb.shape == (dim, )\n",
    "\n",
    "    q = wq[l, :, :] @ xb\n",
    "    if asserts: assert q.shape == (dim, )\n",
    "\n",
    "    k = wk[l, :, :] @ xb\n",
    "    if asserts: assert q.shape == (kv_dim, )\n",
    "\n",
    "    v = wv[l, :, :] @ xb\n",
    "    if asserts: assert q.shape == (kv_dim, )\n",
    "      \n",
    "    # TODO inspect properly\n",
    "    q2 = []\n",
    "    k2 = []\n",
    "\n",
    "    q_tmp = jnp.reshape(q, (dim // 2, 2))\n",
    "    k_tmp = jnp.reshape(k, (dim // 2, 2))\n",
    "\n",
    "    # dim == head_size * n_heads\n",
    "\n",
    "    # Batched gemv\n",
    "    k = jnp.reshape(jnp.einsum('ijk,ik -> ij', toconv2, k_tmp), (dim,))\n",
    "    q = jnp.reshape(jnp.einsum('ijk,ik -> ij', toconv, q_tmp), (dim,))\n",
    "\n",
    "    key_cache_l = key_cache[l, :, :]\n",
    "    key_cache_l = jnp.append(key_cache_l, jnp.reshape(k, (1, dim)), axis=0)\n",
    "    value_cache_l = value_cache[l, :, :]\n",
    "    value_cache_l = jnp.append(value_cache_l, jnp.reshape(v, (1, dim)), axis=0)\n",
    "    keys2.append(key_cache_l)\n",
    "    values2.append(value_cache_l)\n",
    "      \n",
    "    xbs2 = []\n",
    "    for h in range(n_heads):\n",
    "\n",
    "      q2 = q[head_size*h:head_size*(h+1)]\n",
    "      if asserts: assert q2.shape == (head_size,)\n",
    "\n",
    "      # For kv_mul consecutive heads, they share the same kv cache\n",
    "      # reshape key_cache last dim from (kv_dim,) to (kv_mul, head_size)\n",
    "      # generalized einsum reducing the last dim, the rest are batch\n",
    "      att = []\n",
    "\n",
    "      key_index = h // kv_mul\n",
    "        \n",
    "      att = jnp.einsum('ij,j->i', key_cache_l[:, key_index * head_size : (key_index+1) * head_size], q2)\n",
    "\n",
    "      att = att / jnp.sqrt(head_size)\n",
    "\n",
    "      att = softmax(att)\n",
    "        \n",
    "      x_tmp = jnp.einsum('ij,i->j', value_cache_l[:, key_index * head_size : (key_index+1) * head_size], att)\n",
    "\n",
    "      xbs2.append(x_tmp)\n",
    "\n",
    "    # Todo right concat\n",
    "    xb = jnp.concatenate(xbs2, axis=None)\n",
    "\n",
    "    xb2 = wo[l, :, :] @ xb\n",
    "    if asserts: assert xb2.shape == (dim, )\n",
    "\n",
    "    x += xb2\n",
    "\n",
    "    # Rmsnorm and feedforward swiglu\n",
    "\n",
    "    xb = rmsnorm(x, rms_ffn_weight[l, :])\n",
    "\n",
    "    # Now for FFN in PyTorch we have: self.w2(F.silu(self.w1(x)) * self.w3(x))\n",
    "    # first calculate self.w1(x) and self.w3(x)\n",
    "\n",
    "\n",
    "    hb = w1[l, :, :] @ xb\n",
    "    hb2 = w3[l, :, :] @ xb\n",
    "\n",
    "    hb = silu(hb)\n",
    "\n",
    "    hb = hb * hb2\n",
    "\n",
    "\n",
    "    xb = w2[l, :, :] @ hb\n",
    "\n",
    "    x += xb\n",
    "\n",
    "\n",
    "  x = rmsnorm(x, rms_final_weight)\n",
    "  logits = wcls @ x\n",
    "\n",
    "  if asserts: assert logits.shape == (vocab_size,)\n",
    "\n",
    "  for k in keys2:\n",
    "    assert k.shape == (pos+1, kv_dim)\n",
    "  key_cache = jnp.array(keys2)\n",
    "  assert key_cache.shape == (n_layers, pos+1, kv_dim)\n",
    "  value_cache = jnp.array(values2)\n",
    "  assert value_cache.shape == (n_layers, pos+1, kv_dim)\n",
    "  return logits, key_cache, value_cache\n",
    "\n",
    "\n",
    "def str_lookup(string, vocab):\n",
    "    # Find the first perfect match for string in vocab, return its index or -1 if not found\n",
    "    try:\n",
    "        index = vocab.index(string)\n",
    "        return index\n",
    "    except ValueError as err:\n",
    "        return -1\n",
    "\n",
    "def bpe_encode(tokenizer, text):\n",
    "    vocab = tokenizer['vocab']\n",
    "    vocab_scores = tokenizer['vocab_scores']\n",
    "    tokens = []\n",
    "\n",
    "    # First encode every individual character in the input text\n",
    "    for pos, char in enumerate(text):\n",
    "        string = char\n",
    "        id = str_lookup(string, vocab)\n",
    "        if id == -1:\n",
    "            print(f\"not a good prompt at pos {pos}\")\n",
    "            sys.exit(1)\n",
    "        tokens.append(id)\n",
    "\n",
    "    # Merge the best consecutive pair each iteration, according to the scores in vocab_scores\n",
    "    while True:\n",
    "        best_score = -1e10\n",
    "        best_id = -1\n",
    "        best_idx = -1\n",
    "\n",
    "        for i in range(len(tokens) - 1):\n",
    "            # Check if we can merge the pair (tokens[i], tokens[i+1])\n",
    "            # string = vocab[tokens[i]].rstrip(b'\\x00') + vocab[tokens[i + 1]].rstrip(b'\\x00')\n",
    "            string = vocab[tokens[i]] + vocab[tokens[i + 1]]\n",
    "            id = str_lookup(string, vocab)\n",
    "            if id != -1 and vocab_scores[id] > best_score:\n",
    "                # This merge pair exists in vocab! Record its score and position\n",
    "                best_score = vocab_scores[id]\n",
    "                best_id = id\n",
    "                best_idx = i\n",
    "\n",
    "        if best_idx == -1:\n",
    "            break  # We couldn't find any more pairs to merge, so we're done\n",
    "\n",
    "        # Merge the consecutive pair (best_idx, best_idx+1) into new token best_id\n",
    "        tokens[best_idx] = best_id\n",
    "        # Delete token at position best_idx+1, shift the entire sequence back 1\n",
    "        tokens = tokens[0:best_idx + 1] + tokens[best_idx + 2:]\n",
    "\n",
    "    return tokens\n",
    "    \n",
    "def sample(key, logits, temperature, topp):\n",
    "  if temperature == 0:\n",
    "    return jnp.argmax(logits)\n",
    "  else:\n",
    "    logits = logits / temperature\n",
    "    logits = softmax(logits)\n",
    "    if topp <= 0 or topp >= 1:\n",
    "      return jax.random.categorical(key, logits)\n",
    "    else:\n",
    "      raise NotImplementedError(\"not implemented topp\")\n",
    "      # return sample_topp(logits)\n",
    "\n",
    "def run(text, tokenizer, steps, config, weights):\n",
    "  inputs = bpe_encode(tokenizer, text)\n",
    "  key = 0\n",
    "  temperature = 0\n",
    "  topp = 0\n",
    "\n",
    "  n_layers = config['n_layers']\n",
    "  seq_len = config['seq_len']\n",
    "  n_heads = config['n_heads']\n",
    "  dim = config['dim']\n",
    "  n_kv_heads = config['n_kv_heads']\n",
    "  kv_dim = dim // n_heads * n_kv_heads\n",
    "\n",
    "  key_cache = jnp.zeros((n_layers, 0,kv_dim))\n",
    "  value_cache = jnp.zeros((n_layers, 0,kv_dim))\n",
    "\n",
    "  token = 1\n",
    "  for pos in range(steps):\n",
    "    logits, key_cache, value_cache = forward(token, config, weights, key_cache, value_cache)\n",
    "    if pos < len(inputs):\n",
    "      next = inputs[pos]\n",
    "    else:\n",
    "      next = sample(key, logits, temperature, topp)\n",
    "    res = tokenizer['vocab'][next] if pos != -1 else \"<-1>\"\n",
    "    print(\"token # \", pos, \" \", res, \" state \" if (pos < len(inputs)) else \"sampled\", logits)\n",
    "    token = next\n",
    "\n",
    "def loss(token, pos, config, weights, key_cache, value_cache, temperature):\n",
    "  logits = forward(token, config, weights, key_cache, value_cache)\n",
    "  logits = logits / temperature\n",
    "  logits = softmax(logits)\n",
    "  return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4f060329-36b6-40cd-9af7-b4c8a512c2d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('/Users/wmoses/work/llama2.c/stories15M.bin', 'rb') as f:\n",
    "  data = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "ef81f2c9-a6bd-4b13-b05b-7c08303f3b1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "INT_SIZE = 4\n",
    "\n",
    "def load_weights(filename):\n",
    "  config = {}\n",
    "  with open(filename, 'rb') as f:\n",
    "    read_int = lambda: int.from_bytes(f.read(INT_SIZE), byteorder='little')\n",
    "    config['dim'] = read_int()\n",
    "    config['hidden_dim'] = read_int()\n",
    "    config['n_layers'] = read_int()\n",
    "    config['n_heads'] = read_int()\n",
    "    config['n_kv_heads'] = read_int()\n",
    "    config['vocab_size'] = read_int()\n",
    "    config['seq_len'] = read_int()\n",
    "\n",
    "  # Weird encoding with negative size indicating actual weights present.\n",
    "  shared_weights = config['vocab_size'] > 0\n",
    "  config['vocab_size'] = abs(config['vocab_size'])\n",
    "\n",
    "  # Map variables.\n",
    "  dim = config['dim']\n",
    "  hidden_dim = config['hidden_dim']\n",
    "  n_layers = config['n_layers']\n",
    "  n_heads = config['n_heads']\n",
    "  n_kv_heads = config['n_kv_heads']\n",
    "  vocab_size = config['vocab_size']\n",
    "  seq_len = config['seq_len']\n",
    "  print(f\"Loading data with config: {config}\")\n",
    "\n",
    "  # Mmap all data.\n",
    "  config_byte_size = len(config) * INT_SIZE\n",
    "  all_weights = np.memmap(filename, dtype='float32', mode='r', offset=config_byte_size)\n",
    "  offset = 0\n",
    "  weights = {}\n",
    "  def slice_data(name, shape):\n",
    "    nonlocal offset\n",
    "    nonlocal weights\n",
    "    total = np.prod(np.array(shape))\n",
    "    weights[name] = all_weights[offset:offset+total].reshape(shape)\n",
    "    offset += total\n",
    "\n",
    "  assert dim % n_heads == 0\n",
    "  head_size = dim // n_heads\n",
    "\n",
    "  # Take slices of mmaped data in the right order.\n",
    "  slice_data(\"token_embedding_table\",   (vocab_size, dim))\n",
    "  slice_data(\"rms_att_weight\",(n_layers, dim))\n",
    "  slice_data(\"wq\",(n_layers, dim, n_heads * head_size))\n",
    "  slice_data(\"wk\",(n_layers, dim, n_kv_heads * head_size))\n",
    "  slice_data(\"wv\",(n_layers, dim, n_kv_heads * head_size))\n",
    "  slice_data(\"wo\",(n_layers, n_heads * head_size, dim))\n",
    "  slice_data(\"rms_ffn_weight\",(n_layers, dim))\n",
    "  slice_data(\"w1\",(n_layers, hidden_dim, dim))\n",
    "  slice_data(\"w2\",(n_layers, dim, hidden_dim))\n",
    "  slice_data(\"w3\",(n_layers, hidden_dim, dim))\n",
    "  slice_data(\"rms_final_weight\",(dim,))\n",
    "  offset += seq_len * head_size\n",
    "  if shared_weights:\n",
    "    weights[\"wcls\"] = weights[\"token_embedding_table\"]\n",
    "  else:\n",
    "    slice_data(\"wcls\", (vocab_size, dim))\n",
    "  assert offset == len(all_weights), \"haven't read all data\"\n",
    "  return config, weights\n",
    "\n",
    "import struct\n",
    "\n",
    "\n",
    "def read_int_from_file(f):\n",
    "  return int.from_bytes(f.read(INT_SIZE), byteorder='little')\n",
    "\n",
    "def load_tokenizer(filename: str, vocab_size: int):\n",
    "  vocab_scores = []\n",
    "  vocab = []\n",
    "  with open(filename, 'rb') as f:\n",
    "    max_token_length = read_int_from_file(f)\n",
    "    for i in range(vocab_size):\n",
    "      score = struct.unpack('f', f.read(4))[0]\n",
    "      vocab_scores.append(score)\n",
    "      token_length = read_int_from_file(f)\n",
    "      token = f.read(token_length)\n",
    "      if type(token) is not str:\n",
    "        token = token.decode('utf8')\n",
    "      vocab.append(token)\n",
    "  return {'vocab_scores':vocab_scores, 'vocab':vocab, 'max_token_length':max_token_length}\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "2f7ae3c3-ce69-4ada-936f-afbf848138a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading data with config: {'dim': 288, 'hidden_dim': 768, 'n_layers': 6, 'n_heads': 6, 'n_kv_heads': 6, 'vocab_size': 32000, 'seq_len': 256}\n"
     ]
    }
   ],
   "source": [
    "config, weights = load_weights('/Users/wmoses/work/llama2.c/stories15M.bin')\n",
    "\n",
    "tokenizer = load_tokenizer('/Users/wmoses/work/llama2.c/tokenizer.bin', config['vocab_size'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 222,
   "id": "cce297c8-6a8b-4fc8-9aeb-fdcb8d20e1ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "token #  0   D  state  [-6.7907796  0.8281164 -6.790422  ... -6.7907    -6.7906713 -6.790539 ]\n",
      "token #  1   ream  state  [ 0.20367432 -1.5172932   0.20359135 ...  0.20386982  0.20382214\n",
      "  0.2036686 ]\n",
      "token #  2    comes  state  [-6.1911945  0.994138  -6.1906934 ... -6.191182  -6.191168  -6.1909094]\n",
      "token #  3    true  state  [-7.076321   1.8050578 -7.0759583 ... -7.076381  -7.0763884 -7.076169 ]\n",
      "token #  4    this  state  [-10.818716    2.4258273 -10.818541  ... -10.818831  -10.818807\n",
      " -10.818767 ]\n",
      "token #  5    day  state  [-7.903948   -0.22370028 -7.90354    ... -7.904048   -7.9038086\n",
      " -7.903598  ]\n",
      "token #  6   . sampled [-9.333269  2.192757 -9.332975 ... -9.33329  -9.333334 -9.33309 ]\n",
      "token #  7    A sampled [-3.4395585  5.33119   -3.4393303 ... -3.4395003 -3.4394808 -3.439447 ]\n",
      "token #  8    little sampled [-9.060901  -2.7534804 -9.06044   ... -9.060832  -9.060694  -9.060497 ]\n",
      "token #  9    girl sampled [-14.028574   -3.9483452 -14.02855   ... -14.028534  -14.028523\n",
      " -14.028587 ]\n",
      "token #  10    named sampled [-9.035082   2.4355233 -9.03487   ... -9.03515   -9.035015  -9.035023 ]\n",
      "token #  11    Lucy sampled [-8.242906  -0.4827428 -8.242517  ... -8.242999  -8.242893  -8.24273  ]\n",
      "token #  12    is sampled [-6.983634   3.402613  -6.9838886 ... -6.983694  -6.983712  -6.9838705]\n",
      "token #  13    three sampled [-7.8040843  1.1279213 -7.803877  ... -7.804122  -7.8039923 -7.8040333]\n",
      "token #  14    years sampled [-6.9201055  2.9915128 -6.9202075 ... -6.9202685 -6.9202857 -6.9201493]\n",
      "token #  15    old sampled [-7.61737    2.9433894 -7.6176305 ... -7.6175222 -7.617816  -7.617817 ]\n",
      "token #  16   . sampled [-8.47358   3.075387 -8.473835 ... -8.473813 -8.473983 -8.473944]\n",
      "token #  17    She sampled [-3.9300041  7.272192  -3.9298668 ... -3.9300714 -3.9299893 -3.9299312]\n",
      "token #  18    is sampled [-7.625293   4.6335826 -7.625415  ... -7.6252394 -7.6252584 -7.625448 ]\n",
      "token #  19    three sampled [-8.811983   2.8704824 -8.811741  ... -8.811945  -8.811847  -8.811835 ]\n",
      "token #  20    years sampled [-6.3422832  3.50342   -6.3422394 ... -6.342352  -6.342367  -6.342148 ]\n",
      "token #  21    old sampled [-6.8898587  1.6396725 -6.8900547 ... -6.889922  -6.890239  -6.8901772]\n",
      "token #  22   . sampled [-8.903523   3.4991806 -8.903791  ... -8.903733  -8.903912  -8.903855 ]\n",
      "token #  23   <0x0A> sampled [-3.5427623  7.944451  -3.5426369 ... -3.542833  -3.542718  -3.5426257]\n",
      "token #  24   One sampled [-0.3097477   0.6681653  -0.30935907 ... -0.30976772 -0.30958796\n",
      " -0.30936408]\n",
      "token #  25    day sampled [-9.344713   0.3853116 -9.344425  ... -9.344918  -9.344664  -9.344407 ]\n",
      "token #  26   , sampled [-8.514528   2.1890836 -8.514146  ... -8.514616  -8.514551  -8.5141735]\n",
      "token #  27    Lucy sampled [-7.820867   1.6755956 -7.820518  ... -7.8209395 -7.820757  -7.820602 ]\n",
      "token #  28   ' sampled [-5.8109274  3.190032  -5.811026  ... -5.81097   -5.8108797 -5.8110104]\n",
      "token #  29   s sampled [-4.5996804  2.2922804 -4.5997667 ... -4.5998297 -4.599721  -4.599806 ]\n"
     ]
    }
   ],
   "source": [
    "run(\"Dream comes true this day\", tokenizer, 30, config, weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 217,
   "id": "53b7f532-44ba-48ab-8084-b0c8fe8ea39e",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_layers = config['n_layers']\n",
    "seq_len = config['seq_len']\n",
    "n_heads = config['n_heads']\n",
    "dim = config['dim']\n",
    "n_kv_heads = config['n_kv_heads']\n",
    "kv_dim = dim // n_heads * n_kv_heads\n",
    "\n",
    "def partial(func, config):\n",
    "    def sfn(token, weights, key_cache, value_cache):\n",
    "        return func(token, config, weights, key_cache, value_cache)\n",
    "    return sfn\n",
    "\n",
    "pos = 0\n",
    "key_cache = jnp.zeros((n_layers, pos,kv_dim))\n",
    "value_cache = jnp.zeros((n_layers, pos,kv_dim))\n",
    "\n",
    "jfunc = jax.jit(partial(forward, config))\n",
    "mlir = jax.jit(partial(forward, config)).lower(1, weights, key_cache, value_cache).compiler_ir(dialect=\"mhlo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 218,
   "id": "0cc6a65d-3601-4e25-860d-6a23139aa8a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "module @jit_sfn attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {\n",
      "  func.func public @main(%arg0: tensor<i32> {mhlo.sharding = \"{replicated}\"}, %arg1: tensor<6x288xf32> {mhlo.sharding = \"{replicated}\"}, %arg2: tensor<6x288xf32> {mhlo.sharding = \"{replicated}\"}, %arg3: tensor<288xf32> {mhlo.sharding = \"{replicated}\"}, %arg4: tensor<32000x288xf32> {mhlo.sharding = \"{replicated}\"}, %arg5: tensor<6x768x288xf32> {mhlo.sharding = \"{replicated}\"}, %arg6: tensor<6x288x768xf32> {mhlo.sharding = \"{replicated}\"}, %arg7: tensor<6x768x288xf32> {mhlo.sharding = \"{replicated}\"}, %arg8: tensor<32000x288xf32> {mhlo.sharding = \"{replicated}\"}, %arg9: tensor<6x288x288xf32> {mhlo.sharding = \"{replicated}\"}, %arg10: tensor<6x288x288xf32> {mhlo.sharding = \"{replicated}\"}, %arg11: tensor<6x288x288xf32> {mhlo.sharding = \"{replicated}\"}, %arg12: tensor<6x288x288xf32> {mhlo.sharding = \"{replicated}\"}, %arg13: tensor<6x0x288xf32> {mhlo.sharding = \"{replicated}\"}, %arg14: tensor<6x0x288xf32> {mhlo.sharding = \"{replicated}\"}) -> (tensor<32000xf32> {jax.result_info = \"[0]\"}, tensor<6x1x288xf32> {jax.result_info = \"[1]\"}, tensor<6x1x288xf32> {jax.result_info = \"[2]\"}) {\n",
      "    %0 = mhlo.constant dense<0> : tensor<i32>\n",
      "    %1 = mhlo.compare  LT, %arg0, %0,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1>\n",
      "    %2 = mhlo.convert %arg0 : tensor<i32>\n",
      "    %3 = mhlo.constant dense<32000> : tensor<i32>\n",
      "    %4 = mhlo.add %2, %3 : tensor<i32>\n",
      "    %5 = mhlo.select %1, %4, %arg0 : tensor<i1>, tensor<i32>\n",
      "    %6 = mhlo.constant dense<0> : tensor<i32>\n",
      "    %7 = mhlo.constant dense<0> : tensor<i32>\n",
      "    %8 = mhlo.compare  LT, %6, %7,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1>\n",
      "    %9 = mhlo.constant dense<0> : tensor<i32>\n",
      "    %10 = mhlo.constant dense<288> : tensor<i32>\n",
      "    %11 = mhlo.add %9, %10 : tensor<i32>\n",
      "    %12 = mhlo.constant dense<0> : tensor<i32>\n",
      "    %13 = mhlo.select %8, %11, %12 : tensor<i1>, tensor<i32>\n",
      "    %14 = \"mhlo.dynamic_slice\"(%arg4, %5, %13) {slice_sizes = dense<[1, 288]> : tensor<2xi64>} : (tensor<32000x288xf32>, tensor<i32>, tensor<i32>) -> tensor<1x288xf32>\n",
      "    %15 = mhlo.reshape %14 : (tensor<1x288xf32>) -> tensor<288xf32>\n",
      "    %16 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %17 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %18 = mhlo.convert %16 : tensor<f32>\n",
      "    %19 = mhlo.convert %17 : tensor<f32>\n",
      "    %20 = mhlo.power %18, %19 : tensor<f32>\n",
      "    %21 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %22 = mhlo.divide %21, %20 : tensor<f32>\n",
      "    %23 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %24 = mhlo.multiply %23, %22 : tensor<f32>\n",
      "    %25 = mhlo.cosine %24 : tensor<f32>\n",
      "    %26 = mhlo.sine %24 : tensor<f32>\n",
      "    %27 = mhlo.negate %26 : tensor<f32>\n",
      "    %28 = mhlo.convert %25 : tensor<f32>\n",
      "    %29 = mhlo.convert %27 : tensor<f32>\n",
      "    %30 = \"mhlo.broadcast_in_dim\"(%28) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %31 = \"mhlo.broadcast_in_dim\"(%29) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %32 = \"mhlo.concatenate\"(%30, %31) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %33 = mhlo.convert %26 : tensor<f32>\n",
      "    %34 = mhlo.convert %25 : tensor<f32>\n",
      "    %35 = \"mhlo.broadcast_in_dim\"(%33) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %36 = \"mhlo.broadcast_in_dim\"(%34) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %37 = \"mhlo.concatenate\"(%35, %36) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %38 = \"mhlo.broadcast_in_dim\"(%32) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %39 = \"mhlo.broadcast_in_dim\"(%37) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %40 = \"mhlo.concatenate\"(%38, %39) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %41 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %42 = mhlo.constant dense<0.0416666679> : tensor<f32>\n",
      "    %43 = mhlo.convert %41 : tensor<f32>\n",
      "    %44 = mhlo.convert %42 : tensor<f32>\n",
      "    %45 = mhlo.power %43, %44 : tensor<f32>\n",
      "    %46 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %47 = mhlo.divide %46, %45 : tensor<f32>\n",
      "    %48 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %49 = mhlo.multiply %48, %47 : tensor<f32>\n",
      "    %50 = mhlo.cosine %49 : tensor<f32>\n",
      "    %51 = mhlo.sine %49 : tensor<f32>\n",
      "    %52 = mhlo.negate %51 : tensor<f32>\n",
      "    %53 = mhlo.convert %50 : tensor<f32>\n",
      "    %54 = mhlo.convert %52 : tensor<f32>\n",
      "    %55 = \"mhlo.broadcast_in_dim\"(%53) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %56 = \"mhlo.broadcast_in_dim\"(%54) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %57 = \"mhlo.concatenate\"(%55, %56) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %58 = mhlo.convert %51 : tensor<f32>\n",
      "    %59 = mhlo.convert %50 : tensor<f32>\n",
      "    %60 = \"mhlo.broadcast_in_dim\"(%58) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %61 = \"mhlo.broadcast_in_dim\"(%59) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %62 = \"mhlo.concatenate\"(%60, %61) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %63 = \"mhlo.broadcast_in_dim\"(%57) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %64 = \"mhlo.broadcast_in_dim\"(%62) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %65 = \"mhlo.concatenate\"(%63, %64) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %66 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %67 = mhlo.constant dense<0.0833333358> : tensor<f32>\n",
      "    %68 = mhlo.convert %66 : tensor<f32>\n",
      "    %69 = mhlo.convert %67 : tensor<f32>\n",
      "    %70 = mhlo.power %68, %69 : tensor<f32>\n",
      "    %71 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %72 = mhlo.divide %71, %70 : tensor<f32>\n",
      "    %73 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %74 = mhlo.multiply %73, %72 : tensor<f32>\n",
      "    %75 = mhlo.cosine %74 : tensor<f32>\n",
      "    %76 = mhlo.sine %74 : tensor<f32>\n",
      "    %77 = mhlo.negate %76 : tensor<f32>\n",
      "    %78 = mhlo.convert %75 : tensor<f32>\n",
      "    %79 = mhlo.convert %77 : tensor<f32>\n",
      "    %80 = \"mhlo.broadcast_in_dim\"(%78) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %81 = \"mhlo.broadcast_in_dim\"(%79) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %82 = \"mhlo.concatenate\"(%80, %81) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %83 = mhlo.convert %76 : tensor<f32>\n",
      "    %84 = mhlo.convert %75 : tensor<f32>\n",
      "    %85 = \"mhlo.broadcast_in_dim\"(%83) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %86 = \"mhlo.broadcast_in_dim\"(%84) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %87 = \"mhlo.concatenate\"(%85, %86) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %88 = \"mhlo.broadcast_in_dim\"(%82) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %89 = \"mhlo.broadcast_in_dim\"(%87) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %90 = \"mhlo.concatenate\"(%88, %89) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %91 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %92 = mhlo.constant dense<1.250000e-01> : tensor<f32>\n",
      "    %93 = mhlo.convert %91 : tensor<f32>\n",
      "    %94 = mhlo.convert %92 : tensor<f32>\n",
      "    %95 = mhlo.power %93, %94 : tensor<f32>\n",
      "    %96 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %97 = mhlo.divide %96, %95 : tensor<f32>\n",
      "    %98 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %99 = mhlo.multiply %98, %97 : tensor<f32>\n",
      "    %100 = mhlo.cosine %99 : tensor<f32>\n",
      "    %101 = mhlo.sine %99 : tensor<f32>\n",
      "    %102 = mhlo.negate %101 : tensor<f32>\n",
      "    %103 = mhlo.convert %100 : tensor<f32>\n",
      "    %104 = mhlo.convert %102 : tensor<f32>\n",
      "    %105 = \"mhlo.broadcast_in_dim\"(%103) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %106 = \"mhlo.broadcast_in_dim\"(%104) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %107 = \"mhlo.concatenate\"(%105, %106) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %108 = mhlo.convert %101 : tensor<f32>\n",
      "    %109 = mhlo.convert %100 : tensor<f32>\n",
      "    %110 = \"mhlo.broadcast_in_dim\"(%108) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %111 = \"mhlo.broadcast_in_dim\"(%109) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %112 = \"mhlo.concatenate\"(%110, %111) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %113 = \"mhlo.broadcast_in_dim\"(%107) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %114 = \"mhlo.broadcast_in_dim\"(%112) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %115 = \"mhlo.concatenate\"(%113, %114) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %116 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %117 = mhlo.constant dense<0.166666672> : tensor<f32>\n",
      "    %118 = mhlo.convert %116 : tensor<f32>\n",
      "    %119 = mhlo.convert %117 : tensor<f32>\n",
      "    %120 = mhlo.power %118, %119 : tensor<f32>\n",
      "    %121 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %122 = mhlo.divide %121, %120 : tensor<f32>\n",
      "    %123 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %124 = mhlo.multiply %123, %122 : tensor<f32>\n",
      "    %125 = mhlo.cosine %124 : tensor<f32>\n",
      "    %126 = mhlo.sine %124 : tensor<f32>\n",
      "    %127 = mhlo.negate %126 : tensor<f32>\n",
      "    %128 = mhlo.convert %125 : tensor<f32>\n",
      "    %129 = mhlo.convert %127 : tensor<f32>\n",
      "    %130 = \"mhlo.broadcast_in_dim\"(%128) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %131 = \"mhlo.broadcast_in_dim\"(%129) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %132 = \"mhlo.concatenate\"(%130, %131) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %133 = mhlo.convert %126 : tensor<f32>\n",
      "    %134 = mhlo.convert %125 : tensor<f32>\n",
      "    %135 = \"mhlo.broadcast_in_dim\"(%133) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %136 = \"mhlo.broadcast_in_dim\"(%134) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %137 = \"mhlo.concatenate\"(%135, %136) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %138 = \"mhlo.broadcast_in_dim\"(%132) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %139 = \"mhlo.broadcast_in_dim\"(%137) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %140 = \"mhlo.concatenate\"(%138, %139) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %141 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %142 = mhlo.constant dense<0.208333328> : tensor<f32>\n",
      "    %143 = mhlo.convert %141 : tensor<f32>\n",
      "    %144 = mhlo.convert %142 : tensor<f32>\n",
      "    %145 = mhlo.power %143, %144 : tensor<f32>\n",
      "    %146 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %147 = mhlo.divide %146, %145 : tensor<f32>\n",
      "    %148 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %149 = mhlo.multiply %148, %147 : tensor<f32>\n",
      "    %150 = mhlo.cosine %149 : tensor<f32>\n",
      "    %151 = mhlo.sine %149 : tensor<f32>\n",
      "    %152 = mhlo.negate %151 : tensor<f32>\n",
      "    %153 = mhlo.convert %150 : tensor<f32>\n",
      "    %154 = mhlo.convert %152 : tensor<f32>\n",
      "    %155 = \"mhlo.broadcast_in_dim\"(%153) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %156 = \"mhlo.broadcast_in_dim\"(%154) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %157 = \"mhlo.concatenate\"(%155, %156) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %158 = mhlo.convert %151 : tensor<f32>\n",
      "    %159 = mhlo.convert %150 : tensor<f32>\n",
      "    %160 = \"mhlo.broadcast_in_dim\"(%158) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %161 = \"mhlo.broadcast_in_dim\"(%159) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %162 = \"mhlo.concatenate\"(%160, %161) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %163 = \"mhlo.broadcast_in_dim\"(%157) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %164 = \"mhlo.broadcast_in_dim\"(%162) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %165 = \"mhlo.concatenate\"(%163, %164) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %166 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %167 = mhlo.constant dense<2.500000e-01> : tensor<f32>\n",
      "    %168 = mhlo.convert %166 : tensor<f32>\n",
      "    %169 = mhlo.convert %167 : tensor<f32>\n",
      "    %170 = mhlo.power %168, %169 : tensor<f32>\n",
      "    %171 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %172 = mhlo.divide %171, %170 : tensor<f32>\n",
      "    %173 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %174 = mhlo.multiply %173, %172 : tensor<f32>\n",
      "    %175 = mhlo.cosine %174 : tensor<f32>\n",
      "    %176 = mhlo.sine %174 : tensor<f32>\n",
      "    %177 = mhlo.negate %176 : tensor<f32>\n",
      "    %178 = mhlo.convert %175 : tensor<f32>\n",
      "    %179 = mhlo.convert %177 : tensor<f32>\n",
      "    %180 = \"mhlo.broadcast_in_dim\"(%178) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %181 = \"mhlo.broadcast_in_dim\"(%179) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %182 = \"mhlo.concatenate\"(%180, %181) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %183 = mhlo.convert %176 : tensor<f32>\n",
      "    %184 = mhlo.convert %175 : tensor<f32>\n",
      "    %185 = \"mhlo.broadcast_in_dim\"(%183) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %186 = \"mhlo.broadcast_in_dim\"(%184) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %187 = \"mhlo.concatenate\"(%185, %186) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %188 = \"mhlo.broadcast_in_dim\"(%182) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %189 = \"mhlo.broadcast_in_dim\"(%187) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %190 = \"mhlo.concatenate\"(%188, %189) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %191 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %192 = mhlo.constant dense<0.291666657> : tensor<f32>\n",
      "    %193 = mhlo.convert %191 : tensor<f32>\n",
      "    %194 = mhlo.convert %192 : tensor<f32>\n",
      "    %195 = mhlo.power %193, %194 : tensor<f32>\n",
      "    %196 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %197 = mhlo.divide %196, %195 : tensor<f32>\n",
      "    %198 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %199 = mhlo.multiply %198, %197 : tensor<f32>\n",
      "    %200 = mhlo.cosine %199 : tensor<f32>\n",
      "    %201 = mhlo.sine %199 : tensor<f32>\n",
      "    %202 = mhlo.negate %201 : tensor<f32>\n",
      "    %203 = mhlo.convert %200 : tensor<f32>\n",
      "    %204 = mhlo.convert %202 : tensor<f32>\n",
      "    %205 = \"mhlo.broadcast_in_dim\"(%203) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %206 = \"mhlo.broadcast_in_dim\"(%204) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %207 = \"mhlo.concatenate\"(%205, %206) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %208 = mhlo.convert %201 : tensor<f32>\n",
      "    %209 = mhlo.convert %200 : tensor<f32>\n",
      "    %210 = \"mhlo.broadcast_in_dim\"(%208) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %211 = \"mhlo.broadcast_in_dim\"(%209) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %212 = \"mhlo.concatenate\"(%210, %211) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %213 = \"mhlo.broadcast_in_dim\"(%207) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %214 = \"mhlo.broadcast_in_dim\"(%212) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %215 = \"mhlo.concatenate\"(%213, %214) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %216 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %217 = mhlo.constant dense<0.333333343> : tensor<f32>\n",
      "    %218 = mhlo.convert %216 : tensor<f32>\n",
      "    %219 = mhlo.convert %217 : tensor<f32>\n",
      "    %220 = mhlo.power %218, %219 : tensor<f32>\n",
      "    %221 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %222 = mhlo.divide %221, %220 : tensor<f32>\n",
      "    %223 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %224 = mhlo.multiply %223, %222 : tensor<f32>\n",
      "    %225 = mhlo.cosine %224 : tensor<f32>\n",
      "    %226 = mhlo.sine %224 : tensor<f32>\n",
      "    %227 = mhlo.negate %226 : tensor<f32>\n",
      "    %228 = mhlo.convert %225 : tensor<f32>\n",
      "    %229 = mhlo.convert %227 : tensor<f32>\n",
      "    %230 = \"mhlo.broadcast_in_dim\"(%228) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %231 = \"mhlo.broadcast_in_dim\"(%229) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %232 = \"mhlo.concatenate\"(%230, %231) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %233 = mhlo.convert %226 : tensor<f32>\n",
      "    %234 = mhlo.convert %225 : tensor<f32>\n",
      "    %235 = \"mhlo.broadcast_in_dim\"(%233) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %236 = \"mhlo.broadcast_in_dim\"(%234) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %237 = \"mhlo.concatenate\"(%235, %236) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %238 = \"mhlo.broadcast_in_dim\"(%232) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %239 = \"mhlo.broadcast_in_dim\"(%237) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %240 = \"mhlo.concatenate\"(%238, %239) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %241 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %242 = mhlo.constant dense<3.750000e-01> : tensor<f32>\n",
      "    %243 = mhlo.convert %241 : tensor<f32>\n",
      "    %244 = mhlo.convert %242 : tensor<f32>\n",
      "    %245 = mhlo.power %243, %244 : tensor<f32>\n",
      "    %246 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %247 = mhlo.divide %246, %245 : tensor<f32>\n",
      "    %248 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %249 = mhlo.multiply %248, %247 : tensor<f32>\n",
      "    %250 = mhlo.cosine %249 : tensor<f32>\n",
      "    %251 = mhlo.sine %249 : tensor<f32>\n",
      "    %252 = mhlo.negate %251 : tensor<f32>\n",
      "    %253 = mhlo.convert %250 : tensor<f32>\n",
      "    %254 = mhlo.convert %252 : tensor<f32>\n",
      "    %255 = \"mhlo.broadcast_in_dim\"(%253) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %256 = \"mhlo.broadcast_in_dim\"(%254) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %257 = \"mhlo.concatenate\"(%255, %256) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %258 = mhlo.convert %251 : tensor<f32>\n",
      "    %259 = mhlo.convert %250 : tensor<f32>\n",
      "    %260 = \"mhlo.broadcast_in_dim\"(%258) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %261 = \"mhlo.broadcast_in_dim\"(%259) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %262 = \"mhlo.concatenate\"(%260, %261) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %263 = \"mhlo.broadcast_in_dim\"(%257) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %264 = \"mhlo.broadcast_in_dim\"(%262) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %265 = \"mhlo.concatenate\"(%263, %264) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %266 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %267 = mhlo.constant dense<0.416666657> : tensor<f32>\n",
      "    %268 = mhlo.convert %266 : tensor<f32>\n",
      "    %269 = mhlo.convert %267 : tensor<f32>\n",
      "    %270 = mhlo.power %268, %269 : tensor<f32>\n",
      "    %271 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %272 = mhlo.divide %271, %270 : tensor<f32>\n",
      "    %273 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %274 = mhlo.multiply %273, %272 : tensor<f32>\n",
      "    %275 = mhlo.cosine %274 : tensor<f32>\n",
      "    %276 = mhlo.sine %274 : tensor<f32>\n",
      "    %277 = mhlo.negate %276 : tensor<f32>\n",
      "    %278 = mhlo.convert %275 : tensor<f32>\n",
      "    %279 = mhlo.convert %277 : tensor<f32>\n",
      "    %280 = \"mhlo.broadcast_in_dim\"(%278) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %281 = \"mhlo.broadcast_in_dim\"(%279) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %282 = \"mhlo.concatenate\"(%280, %281) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %283 = mhlo.convert %276 : tensor<f32>\n",
      "    %284 = mhlo.convert %275 : tensor<f32>\n",
      "    %285 = \"mhlo.broadcast_in_dim\"(%283) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %286 = \"mhlo.broadcast_in_dim\"(%284) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %287 = \"mhlo.concatenate\"(%285, %286) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %288 = \"mhlo.broadcast_in_dim\"(%282) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %289 = \"mhlo.broadcast_in_dim\"(%287) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %290 = \"mhlo.concatenate\"(%288, %289) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %291 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %292 = mhlo.constant dense<0.458333343> : tensor<f32>\n",
      "    %293 = mhlo.convert %291 : tensor<f32>\n",
      "    %294 = mhlo.convert %292 : tensor<f32>\n",
      "    %295 = mhlo.power %293, %294 : tensor<f32>\n",
      "    %296 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %297 = mhlo.divide %296, %295 : tensor<f32>\n",
      "    %298 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %299 = mhlo.multiply %298, %297 : tensor<f32>\n",
      "    %300 = mhlo.cosine %299 : tensor<f32>\n",
      "    %301 = mhlo.sine %299 : tensor<f32>\n",
      "    %302 = mhlo.negate %301 : tensor<f32>\n",
      "    %303 = mhlo.convert %300 : tensor<f32>\n",
      "    %304 = mhlo.convert %302 : tensor<f32>\n",
      "    %305 = \"mhlo.broadcast_in_dim\"(%303) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %306 = \"mhlo.broadcast_in_dim\"(%304) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %307 = \"mhlo.concatenate\"(%305, %306) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %308 = mhlo.convert %301 : tensor<f32>\n",
      "    %309 = mhlo.convert %300 : tensor<f32>\n",
      "    %310 = \"mhlo.broadcast_in_dim\"(%308) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %311 = \"mhlo.broadcast_in_dim\"(%309) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %312 = \"mhlo.concatenate\"(%310, %311) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %313 = \"mhlo.broadcast_in_dim\"(%307) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %314 = \"mhlo.broadcast_in_dim\"(%312) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %315 = \"mhlo.concatenate\"(%313, %314) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %316 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %317 = mhlo.constant dense<5.000000e-01> : tensor<f32>\n",
      "    %318 = mhlo.convert %316 : tensor<f32>\n",
      "    %319 = mhlo.convert %317 : tensor<f32>\n",
      "    %320 = mhlo.power %318, %319 : tensor<f32>\n",
      "    %321 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %322 = mhlo.divide %321, %320 : tensor<f32>\n",
      "    %323 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %324 = mhlo.multiply %323, %322 : tensor<f32>\n",
      "    %325 = mhlo.cosine %324 : tensor<f32>\n",
      "    %326 = mhlo.sine %324 : tensor<f32>\n",
      "    %327 = mhlo.negate %326 : tensor<f32>\n",
      "    %328 = mhlo.convert %325 : tensor<f32>\n",
      "    %329 = mhlo.convert %327 : tensor<f32>\n",
      "    %330 = \"mhlo.broadcast_in_dim\"(%328) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %331 = \"mhlo.broadcast_in_dim\"(%329) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %332 = \"mhlo.concatenate\"(%330, %331) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %333 = mhlo.convert %326 : tensor<f32>\n",
      "    %334 = mhlo.convert %325 : tensor<f32>\n",
      "    %335 = \"mhlo.broadcast_in_dim\"(%333) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %336 = \"mhlo.broadcast_in_dim\"(%334) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %337 = \"mhlo.concatenate\"(%335, %336) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %338 = \"mhlo.broadcast_in_dim\"(%332) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %339 = \"mhlo.broadcast_in_dim\"(%337) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %340 = \"mhlo.concatenate\"(%338, %339) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %341 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %342 = mhlo.constant dense<0.541666687> : tensor<f32>\n",
      "    %343 = mhlo.convert %341 : tensor<f32>\n",
      "    %344 = mhlo.convert %342 : tensor<f32>\n",
      "    %345 = mhlo.power %343, %344 : tensor<f32>\n",
      "    %346 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %347 = mhlo.divide %346, %345 : tensor<f32>\n",
      "    %348 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %349 = mhlo.multiply %348, %347 : tensor<f32>\n",
      "    %350 = mhlo.cosine %349 : tensor<f32>\n",
      "    %351 = mhlo.sine %349 : tensor<f32>\n",
      "    %352 = mhlo.negate %351 : tensor<f32>\n",
      "    %353 = mhlo.convert %350 : tensor<f32>\n",
      "    %354 = mhlo.convert %352 : tensor<f32>\n",
      "    %355 = \"mhlo.broadcast_in_dim\"(%353) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %356 = \"mhlo.broadcast_in_dim\"(%354) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %357 = \"mhlo.concatenate\"(%355, %356) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %358 = mhlo.convert %351 : tensor<f32>\n",
      "    %359 = mhlo.convert %350 : tensor<f32>\n",
      "    %360 = \"mhlo.broadcast_in_dim\"(%358) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %361 = \"mhlo.broadcast_in_dim\"(%359) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %362 = \"mhlo.concatenate\"(%360, %361) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %363 = \"mhlo.broadcast_in_dim\"(%357) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %364 = \"mhlo.broadcast_in_dim\"(%362) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %365 = \"mhlo.concatenate\"(%363, %364) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %366 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %367 = mhlo.constant dense<0.583333313> : tensor<f32>\n",
      "    %368 = mhlo.convert %366 : tensor<f32>\n",
      "    %369 = mhlo.convert %367 : tensor<f32>\n",
      "    %370 = mhlo.power %368, %369 : tensor<f32>\n",
      "    %371 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %372 = mhlo.divide %371, %370 : tensor<f32>\n",
      "    %373 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %374 = mhlo.multiply %373, %372 : tensor<f32>\n",
      "    %375 = mhlo.cosine %374 : tensor<f32>\n",
      "    %376 = mhlo.sine %374 : tensor<f32>\n",
      "    %377 = mhlo.negate %376 : tensor<f32>\n",
      "    %378 = mhlo.convert %375 : tensor<f32>\n",
      "    %379 = mhlo.convert %377 : tensor<f32>\n",
      "    %380 = \"mhlo.broadcast_in_dim\"(%378) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %381 = \"mhlo.broadcast_in_dim\"(%379) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %382 = \"mhlo.concatenate\"(%380, %381) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %383 = mhlo.convert %376 : tensor<f32>\n",
      "    %384 = mhlo.convert %375 : tensor<f32>\n",
      "    %385 = \"mhlo.broadcast_in_dim\"(%383) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %386 = \"mhlo.broadcast_in_dim\"(%384) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %387 = \"mhlo.concatenate\"(%385, %386) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %388 = \"mhlo.broadcast_in_dim\"(%382) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %389 = \"mhlo.broadcast_in_dim\"(%387) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %390 = \"mhlo.concatenate\"(%388, %389) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %391 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %392 = mhlo.constant dense<6.250000e-01> : tensor<f32>\n",
      "    %393 = mhlo.convert %391 : tensor<f32>\n",
      "    %394 = mhlo.convert %392 : tensor<f32>\n",
      "    %395 = mhlo.power %393, %394 : tensor<f32>\n",
      "    %396 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %397 = mhlo.divide %396, %395 : tensor<f32>\n",
      "    %398 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %399 = mhlo.multiply %398, %397 : tensor<f32>\n",
      "    %400 = mhlo.cosine %399 : tensor<f32>\n",
      "    %401 = mhlo.sine %399 : tensor<f32>\n",
      "    %402 = mhlo.negate %401 : tensor<f32>\n",
      "    %403 = mhlo.convert %400 : tensor<f32>\n",
      "    %404 = mhlo.convert %402 : tensor<f32>\n",
      "    %405 = \"mhlo.broadcast_in_dim\"(%403) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %406 = \"mhlo.broadcast_in_dim\"(%404) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %407 = \"mhlo.concatenate\"(%405, %406) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %408 = mhlo.convert %401 : tensor<f32>\n",
      "    %409 = mhlo.convert %400 : tensor<f32>\n",
      "    %410 = \"mhlo.broadcast_in_dim\"(%408) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %411 = \"mhlo.broadcast_in_dim\"(%409) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %412 = \"mhlo.concatenate\"(%410, %411) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %413 = \"mhlo.broadcast_in_dim\"(%407) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %414 = \"mhlo.broadcast_in_dim\"(%412) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %415 = \"mhlo.concatenate\"(%413, %414) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %416 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %417 = mhlo.constant dense<0.666666686> : tensor<f32>\n",
      "    %418 = mhlo.convert %416 : tensor<f32>\n",
      "    %419 = mhlo.convert %417 : tensor<f32>\n",
      "    %420 = mhlo.power %418, %419 : tensor<f32>\n",
      "    %421 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %422 = mhlo.divide %421, %420 : tensor<f32>\n",
      "    %423 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %424 = mhlo.multiply %423, %422 : tensor<f32>\n",
      "    %425 = mhlo.cosine %424 : tensor<f32>\n",
      "    %426 = mhlo.sine %424 : tensor<f32>\n",
      "    %427 = mhlo.negate %426 : tensor<f32>\n",
      "    %428 = mhlo.convert %425 : tensor<f32>\n",
      "    %429 = mhlo.convert %427 : tensor<f32>\n",
      "    %430 = \"mhlo.broadcast_in_dim\"(%428) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %431 = \"mhlo.broadcast_in_dim\"(%429) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %432 = \"mhlo.concatenate\"(%430, %431) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %433 = mhlo.convert %426 : tensor<f32>\n",
      "    %434 = mhlo.convert %425 : tensor<f32>\n",
      "    %435 = \"mhlo.broadcast_in_dim\"(%433) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %436 = \"mhlo.broadcast_in_dim\"(%434) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %437 = \"mhlo.concatenate\"(%435, %436) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %438 = \"mhlo.broadcast_in_dim\"(%432) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %439 = \"mhlo.broadcast_in_dim\"(%437) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %440 = \"mhlo.concatenate\"(%438, %439) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %441 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %442 = mhlo.constant dense<0.708333313> : tensor<f32>\n",
      "    %443 = mhlo.convert %441 : tensor<f32>\n",
      "    %444 = mhlo.convert %442 : tensor<f32>\n",
      "    %445 = mhlo.power %443, %444 : tensor<f32>\n",
      "    %446 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %447 = mhlo.divide %446, %445 : tensor<f32>\n",
      "    %448 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %449 = mhlo.multiply %448, %447 : tensor<f32>\n",
      "    %450 = mhlo.cosine %449 : tensor<f32>\n",
      "    %451 = mhlo.sine %449 : tensor<f32>\n",
      "    %452 = mhlo.negate %451 : tensor<f32>\n",
      "    %453 = mhlo.convert %450 : tensor<f32>\n",
      "    %454 = mhlo.convert %452 : tensor<f32>\n",
      "    %455 = \"mhlo.broadcast_in_dim\"(%453) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %456 = \"mhlo.broadcast_in_dim\"(%454) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %457 = \"mhlo.concatenate\"(%455, %456) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %458 = mhlo.convert %451 : tensor<f32>\n",
      "    %459 = mhlo.convert %450 : tensor<f32>\n",
      "    %460 = \"mhlo.broadcast_in_dim\"(%458) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %461 = \"mhlo.broadcast_in_dim\"(%459) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %462 = \"mhlo.concatenate\"(%460, %461) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %463 = \"mhlo.broadcast_in_dim\"(%457) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %464 = \"mhlo.broadcast_in_dim\"(%462) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %465 = \"mhlo.concatenate\"(%463, %464) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %466 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %467 = mhlo.constant dense<7.500000e-01> : tensor<f32>\n",
      "    %468 = mhlo.convert %466 : tensor<f32>\n",
      "    %469 = mhlo.convert %467 : tensor<f32>\n",
      "    %470 = mhlo.power %468, %469 : tensor<f32>\n",
      "    %471 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %472 = mhlo.divide %471, %470 : tensor<f32>\n",
      "    %473 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %474 = mhlo.multiply %473, %472 : tensor<f32>\n",
      "    %475 = mhlo.cosine %474 : tensor<f32>\n",
      "    %476 = mhlo.sine %474 : tensor<f32>\n",
      "    %477 = mhlo.negate %476 : tensor<f32>\n",
      "    %478 = mhlo.convert %475 : tensor<f32>\n",
      "    %479 = mhlo.convert %477 : tensor<f32>\n",
      "    %480 = \"mhlo.broadcast_in_dim\"(%478) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %481 = \"mhlo.broadcast_in_dim\"(%479) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %482 = \"mhlo.concatenate\"(%480, %481) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %483 = mhlo.convert %476 : tensor<f32>\n",
      "    %484 = mhlo.convert %475 : tensor<f32>\n",
      "    %485 = \"mhlo.broadcast_in_dim\"(%483) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %486 = \"mhlo.broadcast_in_dim\"(%484) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %487 = \"mhlo.concatenate\"(%485, %486) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %488 = \"mhlo.broadcast_in_dim\"(%482) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %489 = \"mhlo.broadcast_in_dim\"(%487) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %490 = \"mhlo.concatenate\"(%488, %489) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %491 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %492 = mhlo.constant dense<0.791666686> : tensor<f32>\n",
      "    %493 = mhlo.convert %491 : tensor<f32>\n",
      "    %494 = mhlo.convert %492 : tensor<f32>\n",
      "    %495 = mhlo.power %493, %494 : tensor<f32>\n",
      "    %496 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %497 = mhlo.divide %496, %495 : tensor<f32>\n",
      "    %498 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %499 = mhlo.multiply %498, %497 : tensor<f32>\n",
      "    %500 = mhlo.cosine %499 : tensor<f32>\n",
      "    %501 = mhlo.sine %499 : tensor<f32>\n",
      "    %502 = mhlo.negate %501 : tensor<f32>\n",
      "    %503 = mhlo.convert %500 : tensor<f32>\n",
      "    %504 = mhlo.convert %502 : tensor<f32>\n",
      "    %505 = \"mhlo.broadcast_in_dim\"(%503) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %506 = \"mhlo.broadcast_in_dim\"(%504) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %507 = \"mhlo.concatenate\"(%505, %506) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %508 = mhlo.convert %501 : tensor<f32>\n",
      "    %509 = mhlo.convert %500 : tensor<f32>\n",
      "    %510 = \"mhlo.broadcast_in_dim\"(%508) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %511 = \"mhlo.broadcast_in_dim\"(%509) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %512 = \"mhlo.concatenate\"(%510, %511) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %513 = \"mhlo.broadcast_in_dim\"(%507) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %514 = \"mhlo.broadcast_in_dim\"(%512) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %515 = \"mhlo.concatenate\"(%513, %514) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %516 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %517 = mhlo.constant dense<0.833333313> : tensor<f32>\n",
      "    %518 = mhlo.convert %516 : tensor<f32>\n",
      "    %519 = mhlo.convert %517 : tensor<f32>\n",
      "    %520 = mhlo.power %518, %519 : tensor<f32>\n",
      "    %521 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %522 = mhlo.divide %521, %520 : tensor<f32>\n",
      "    %523 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %524 = mhlo.multiply %523, %522 : tensor<f32>\n",
      "    %525 = mhlo.cosine %524 : tensor<f32>\n",
      "    %526 = mhlo.sine %524 : tensor<f32>\n",
      "    %527 = mhlo.negate %526 : tensor<f32>\n",
      "    %528 = mhlo.convert %525 : tensor<f32>\n",
      "    %529 = mhlo.convert %527 : tensor<f32>\n",
      "    %530 = \"mhlo.broadcast_in_dim\"(%528) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %531 = \"mhlo.broadcast_in_dim\"(%529) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %532 = \"mhlo.concatenate\"(%530, %531) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %533 = mhlo.convert %526 : tensor<f32>\n",
      "    %534 = mhlo.convert %525 : tensor<f32>\n",
      "    %535 = \"mhlo.broadcast_in_dim\"(%533) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %536 = \"mhlo.broadcast_in_dim\"(%534) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %537 = \"mhlo.concatenate\"(%535, %536) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %538 = \"mhlo.broadcast_in_dim\"(%532) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %539 = \"mhlo.broadcast_in_dim\"(%537) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %540 = \"mhlo.concatenate\"(%538, %539) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %541 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %542 = mhlo.constant dense<8.750000e-01> : tensor<f32>\n",
      "    %543 = mhlo.convert %541 : tensor<f32>\n",
      "    %544 = mhlo.convert %542 : tensor<f32>\n",
      "    %545 = mhlo.power %543, %544 : tensor<f32>\n",
      "    %546 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %547 = mhlo.divide %546, %545 : tensor<f32>\n",
      "    %548 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %549 = mhlo.multiply %548, %547 : tensor<f32>\n",
      "    %550 = mhlo.cosine %549 : tensor<f32>\n",
      "    %551 = mhlo.sine %549 : tensor<f32>\n",
      "    %552 = mhlo.negate %551 : tensor<f32>\n",
      "    %553 = mhlo.convert %550 : tensor<f32>\n",
      "    %554 = mhlo.convert %552 : tensor<f32>\n",
      "    %555 = \"mhlo.broadcast_in_dim\"(%553) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %556 = \"mhlo.broadcast_in_dim\"(%554) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %557 = \"mhlo.concatenate\"(%555, %556) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %558 = mhlo.convert %551 : tensor<f32>\n",
      "    %559 = mhlo.convert %550 : tensor<f32>\n",
      "    %560 = \"mhlo.broadcast_in_dim\"(%558) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %561 = \"mhlo.broadcast_in_dim\"(%559) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %562 = \"mhlo.concatenate\"(%560, %561) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %563 = \"mhlo.broadcast_in_dim\"(%557) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %564 = \"mhlo.broadcast_in_dim\"(%562) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %565 = \"mhlo.concatenate\"(%563, %564) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %566 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %567 = mhlo.constant dense<0.916666686> : tensor<f32>\n",
      "    %568 = mhlo.convert %566 : tensor<f32>\n",
      "    %569 = mhlo.convert %567 : tensor<f32>\n",
      "    %570 = mhlo.power %568, %569 : tensor<f32>\n",
      "    %571 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %572 = mhlo.divide %571, %570 : tensor<f32>\n",
      "    %573 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %574 = mhlo.multiply %573, %572 : tensor<f32>\n",
      "    %575 = mhlo.cosine %574 : tensor<f32>\n",
      "    %576 = mhlo.sine %574 : tensor<f32>\n",
      "    %577 = mhlo.negate %576 : tensor<f32>\n",
      "    %578 = mhlo.convert %575 : tensor<f32>\n",
      "    %579 = mhlo.convert %577 : tensor<f32>\n",
      "    %580 = \"mhlo.broadcast_in_dim\"(%578) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %581 = \"mhlo.broadcast_in_dim\"(%579) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %582 = \"mhlo.concatenate\"(%580, %581) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %583 = mhlo.convert %576 : tensor<f32>\n",
      "    %584 = mhlo.convert %575 : tensor<f32>\n",
      "    %585 = \"mhlo.broadcast_in_dim\"(%583) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %586 = \"mhlo.broadcast_in_dim\"(%584) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %587 = \"mhlo.concatenate\"(%585, %586) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %588 = \"mhlo.broadcast_in_dim\"(%582) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %589 = \"mhlo.broadcast_in_dim\"(%587) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %590 = \"mhlo.concatenate\"(%588, %589) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %591 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %592 = mhlo.constant dense<0.958333313> : tensor<f32>\n",
      "    %593 = mhlo.convert %591 : tensor<f32>\n",
      "    %594 = mhlo.convert %592 : tensor<f32>\n",
      "    %595 = mhlo.power %593, %594 : tensor<f32>\n",
      "    %596 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %597 = mhlo.divide %596, %595 : tensor<f32>\n",
      "    %598 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %599 = mhlo.multiply %598, %597 : tensor<f32>\n",
      "    %600 = mhlo.cosine %599 : tensor<f32>\n",
      "    %601 = mhlo.sine %599 : tensor<f32>\n",
      "    %602 = mhlo.negate %601 : tensor<f32>\n",
      "    %603 = mhlo.convert %600 : tensor<f32>\n",
      "    %604 = mhlo.convert %602 : tensor<f32>\n",
      "    %605 = \"mhlo.broadcast_in_dim\"(%603) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %606 = \"mhlo.broadcast_in_dim\"(%604) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %607 = \"mhlo.concatenate\"(%605, %606) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %608 = mhlo.convert %601 : tensor<f32>\n",
      "    %609 = mhlo.convert %600 : tensor<f32>\n",
      "    %610 = \"mhlo.broadcast_in_dim\"(%608) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %611 = \"mhlo.broadcast_in_dim\"(%609) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %612 = \"mhlo.concatenate\"(%610, %611) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %613 = \"mhlo.broadcast_in_dim\"(%607) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %614 = \"mhlo.broadcast_in_dim\"(%612) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %615 = \"mhlo.concatenate\"(%613, %614) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %616 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %617 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %618 = mhlo.convert %616 : tensor<f32>\n",
      "    %619 = mhlo.convert %617 : tensor<f32>\n",
      "    %620 = mhlo.power %618, %619 : tensor<f32>\n",
      "    %621 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %622 = mhlo.divide %621, %620 : tensor<f32>\n",
      "    %623 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %624 = mhlo.multiply %623, %622 : tensor<f32>\n",
      "    %625 = mhlo.cosine %624 : tensor<f32>\n",
      "    %626 = mhlo.sine %624 : tensor<f32>\n",
      "    %627 = mhlo.negate %626 : tensor<f32>\n",
      "    %628 = mhlo.convert %625 : tensor<f32>\n",
      "    %629 = mhlo.convert %627 : tensor<f32>\n",
      "    %630 = \"mhlo.broadcast_in_dim\"(%628) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %631 = \"mhlo.broadcast_in_dim\"(%629) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %632 = \"mhlo.concatenate\"(%630, %631) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %633 = mhlo.convert %626 : tensor<f32>\n",
      "    %634 = mhlo.convert %625 : tensor<f32>\n",
      "    %635 = \"mhlo.broadcast_in_dim\"(%633) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %636 = \"mhlo.broadcast_in_dim\"(%634) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %637 = \"mhlo.concatenate\"(%635, %636) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %638 = \"mhlo.broadcast_in_dim\"(%632) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %639 = \"mhlo.broadcast_in_dim\"(%637) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %640 = \"mhlo.concatenate\"(%638, %639) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %641 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %642 = mhlo.constant dense<0.0416666679> : tensor<f32>\n",
      "    %643 = mhlo.convert %641 : tensor<f32>\n",
      "    %644 = mhlo.convert %642 : tensor<f32>\n",
      "    %645 = mhlo.power %643, %644 : tensor<f32>\n",
      "    %646 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %647 = mhlo.divide %646, %645 : tensor<f32>\n",
      "    %648 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %649 = mhlo.multiply %648, %647 : tensor<f32>\n",
      "    %650 = mhlo.cosine %649 : tensor<f32>\n",
      "    %651 = mhlo.sine %649 : tensor<f32>\n",
      "    %652 = mhlo.negate %651 : tensor<f32>\n",
      "    %653 = mhlo.convert %650 : tensor<f32>\n",
      "    %654 = mhlo.convert %652 : tensor<f32>\n",
      "    %655 = \"mhlo.broadcast_in_dim\"(%653) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %656 = \"mhlo.broadcast_in_dim\"(%654) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %657 = \"mhlo.concatenate\"(%655, %656) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %658 = mhlo.convert %651 : tensor<f32>\n",
      "    %659 = mhlo.convert %650 : tensor<f32>\n",
      "    %660 = \"mhlo.broadcast_in_dim\"(%658) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %661 = \"mhlo.broadcast_in_dim\"(%659) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %662 = \"mhlo.concatenate\"(%660, %661) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %663 = \"mhlo.broadcast_in_dim\"(%657) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %664 = \"mhlo.broadcast_in_dim\"(%662) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %665 = \"mhlo.concatenate\"(%663, %664) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %666 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %667 = mhlo.constant dense<0.0833333358> : tensor<f32>\n",
      "    %668 = mhlo.convert %666 : tensor<f32>\n",
      "    %669 = mhlo.convert %667 : tensor<f32>\n",
      "    %670 = mhlo.power %668, %669 : tensor<f32>\n",
      "    %671 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %672 = mhlo.divide %671, %670 : tensor<f32>\n",
      "    %673 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %674 = mhlo.multiply %673, %672 : tensor<f32>\n",
      "    %675 = mhlo.cosine %674 : tensor<f32>\n",
      "    %676 = mhlo.sine %674 : tensor<f32>\n",
      "    %677 = mhlo.negate %676 : tensor<f32>\n",
      "    %678 = mhlo.convert %675 : tensor<f32>\n",
      "    %679 = mhlo.convert %677 : tensor<f32>\n",
      "    %680 = \"mhlo.broadcast_in_dim\"(%678) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %681 = \"mhlo.broadcast_in_dim\"(%679) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %682 = \"mhlo.concatenate\"(%680, %681) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %683 = mhlo.convert %676 : tensor<f32>\n",
      "    %684 = mhlo.convert %675 : tensor<f32>\n",
      "    %685 = \"mhlo.broadcast_in_dim\"(%683) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %686 = \"mhlo.broadcast_in_dim\"(%684) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %687 = \"mhlo.concatenate\"(%685, %686) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %688 = \"mhlo.broadcast_in_dim\"(%682) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %689 = \"mhlo.broadcast_in_dim\"(%687) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %690 = \"mhlo.concatenate\"(%688, %689) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %691 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %692 = mhlo.constant dense<1.250000e-01> : tensor<f32>\n",
      "    %693 = mhlo.convert %691 : tensor<f32>\n",
      "    %694 = mhlo.convert %692 : tensor<f32>\n",
      "    %695 = mhlo.power %693, %694 : tensor<f32>\n",
      "    %696 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %697 = mhlo.divide %696, %695 : tensor<f32>\n",
      "    %698 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %699 = mhlo.multiply %698, %697 : tensor<f32>\n",
      "    %700 = mhlo.cosine %699 : tensor<f32>\n",
      "    %701 = mhlo.sine %699 : tensor<f32>\n",
      "    %702 = mhlo.negate %701 : tensor<f32>\n",
      "    %703 = mhlo.convert %700 : tensor<f32>\n",
      "    %704 = mhlo.convert %702 : tensor<f32>\n",
      "    %705 = \"mhlo.broadcast_in_dim\"(%703) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %706 = \"mhlo.broadcast_in_dim\"(%704) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %707 = \"mhlo.concatenate\"(%705, %706) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %708 = mhlo.convert %701 : tensor<f32>\n",
      "    %709 = mhlo.convert %700 : tensor<f32>\n",
      "    %710 = \"mhlo.broadcast_in_dim\"(%708) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %711 = \"mhlo.broadcast_in_dim\"(%709) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %712 = \"mhlo.concatenate\"(%710, %711) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %713 = \"mhlo.broadcast_in_dim\"(%707) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %714 = \"mhlo.broadcast_in_dim\"(%712) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %715 = \"mhlo.concatenate\"(%713, %714) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %716 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %717 = mhlo.constant dense<0.166666672> : tensor<f32>\n",
      "    %718 = mhlo.convert %716 : tensor<f32>\n",
      "    %719 = mhlo.convert %717 : tensor<f32>\n",
      "    %720 = mhlo.power %718, %719 : tensor<f32>\n",
      "    %721 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %722 = mhlo.divide %721, %720 : tensor<f32>\n",
      "    %723 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %724 = mhlo.multiply %723, %722 : tensor<f32>\n",
      "    %725 = mhlo.cosine %724 : tensor<f32>\n",
      "    %726 = mhlo.sine %724 : tensor<f32>\n",
      "    %727 = mhlo.negate %726 : tensor<f32>\n",
      "    %728 = mhlo.convert %725 : tensor<f32>\n",
      "    %729 = mhlo.convert %727 : tensor<f32>\n",
      "    %730 = \"mhlo.broadcast_in_dim\"(%728) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %731 = \"mhlo.broadcast_in_dim\"(%729) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %732 = \"mhlo.concatenate\"(%730, %731) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %733 = mhlo.convert %726 : tensor<f32>\n",
      "    %734 = mhlo.convert %725 : tensor<f32>\n",
      "    %735 = \"mhlo.broadcast_in_dim\"(%733) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %736 = \"mhlo.broadcast_in_dim\"(%734) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %737 = \"mhlo.concatenate\"(%735, %736) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %738 = \"mhlo.broadcast_in_dim\"(%732) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %739 = \"mhlo.broadcast_in_dim\"(%737) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %740 = \"mhlo.concatenate\"(%738, %739) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %741 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %742 = mhlo.constant dense<0.208333328> : tensor<f32>\n",
      "    %743 = mhlo.convert %741 : tensor<f32>\n",
      "    %744 = mhlo.convert %742 : tensor<f32>\n",
      "    %745 = mhlo.power %743, %744 : tensor<f32>\n",
      "    %746 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %747 = mhlo.divide %746, %745 : tensor<f32>\n",
      "    %748 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %749 = mhlo.multiply %748, %747 : tensor<f32>\n",
      "    %750 = mhlo.cosine %749 : tensor<f32>\n",
      "    %751 = mhlo.sine %749 : tensor<f32>\n",
      "    %752 = mhlo.negate %751 : tensor<f32>\n",
      "    %753 = mhlo.convert %750 : tensor<f32>\n",
      "    %754 = mhlo.convert %752 : tensor<f32>\n",
      "    %755 = \"mhlo.broadcast_in_dim\"(%753) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %756 = \"mhlo.broadcast_in_dim\"(%754) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %757 = \"mhlo.concatenate\"(%755, %756) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %758 = mhlo.convert %751 : tensor<f32>\n",
      "    %759 = mhlo.convert %750 : tensor<f32>\n",
      "    %760 = \"mhlo.broadcast_in_dim\"(%758) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %761 = \"mhlo.broadcast_in_dim\"(%759) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %762 = \"mhlo.concatenate\"(%760, %761) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %763 = \"mhlo.broadcast_in_dim\"(%757) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %764 = \"mhlo.broadcast_in_dim\"(%762) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %765 = \"mhlo.concatenate\"(%763, %764) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %766 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %767 = mhlo.constant dense<2.500000e-01> : tensor<f32>\n",
      "    %768 = mhlo.convert %766 : tensor<f32>\n",
      "    %769 = mhlo.convert %767 : tensor<f32>\n",
      "    %770 = mhlo.power %768, %769 : tensor<f32>\n",
      "    %771 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %772 = mhlo.divide %771, %770 : tensor<f32>\n",
      "    %773 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %774 = mhlo.multiply %773, %772 : tensor<f32>\n",
      "    %775 = mhlo.cosine %774 : tensor<f32>\n",
      "    %776 = mhlo.sine %774 : tensor<f32>\n",
      "    %777 = mhlo.negate %776 : tensor<f32>\n",
      "    %778 = mhlo.convert %775 : tensor<f32>\n",
      "    %779 = mhlo.convert %777 : tensor<f32>\n",
      "    %780 = \"mhlo.broadcast_in_dim\"(%778) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %781 = \"mhlo.broadcast_in_dim\"(%779) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %782 = \"mhlo.concatenate\"(%780, %781) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %783 = mhlo.convert %776 : tensor<f32>\n",
      "    %784 = mhlo.convert %775 : tensor<f32>\n",
      "    %785 = \"mhlo.broadcast_in_dim\"(%783) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %786 = \"mhlo.broadcast_in_dim\"(%784) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %787 = \"mhlo.concatenate\"(%785, %786) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %788 = \"mhlo.broadcast_in_dim\"(%782) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %789 = \"mhlo.broadcast_in_dim\"(%787) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %790 = \"mhlo.concatenate\"(%788, %789) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %791 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %792 = mhlo.constant dense<0.291666657> : tensor<f32>\n",
      "    %793 = mhlo.convert %791 : tensor<f32>\n",
      "    %794 = mhlo.convert %792 : tensor<f32>\n",
      "    %795 = mhlo.power %793, %794 : tensor<f32>\n",
      "    %796 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %797 = mhlo.divide %796, %795 : tensor<f32>\n",
      "    %798 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %799 = mhlo.multiply %798, %797 : tensor<f32>\n",
      "    %800 = mhlo.cosine %799 : tensor<f32>\n",
      "    %801 = mhlo.sine %799 : tensor<f32>\n",
      "    %802 = mhlo.negate %801 : tensor<f32>\n",
      "    %803 = mhlo.convert %800 : tensor<f32>\n",
      "    %804 = mhlo.convert %802 : tensor<f32>\n",
      "    %805 = \"mhlo.broadcast_in_dim\"(%803) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %806 = \"mhlo.broadcast_in_dim\"(%804) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %807 = \"mhlo.concatenate\"(%805, %806) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %808 = mhlo.convert %801 : tensor<f32>\n",
      "    %809 = mhlo.convert %800 : tensor<f32>\n",
      "    %810 = \"mhlo.broadcast_in_dim\"(%808) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %811 = \"mhlo.broadcast_in_dim\"(%809) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %812 = \"mhlo.concatenate\"(%810, %811) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %813 = \"mhlo.broadcast_in_dim\"(%807) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %814 = \"mhlo.broadcast_in_dim\"(%812) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %815 = \"mhlo.concatenate\"(%813, %814) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %816 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %817 = mhlo.constant dense<0.333333343> : tensor<f32>\n",
      "    %818 = mhlo.convert %816 : tensor<f32>\n",
      "    %819 = mhlo.convert %817 : tensor<f32>\n",
      "    %820 = mhlo.power %818, %819 : tensor<f32>\n",
      "    %821 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %822 = mhlo.divide %821, %820 : tensor<f32>\n",
      "    %823 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %824 = mhlo.multiply %823, %822 : tensor<f32>\n",
      "    %825 = mhlo.cosine %824 : tensor<f32>\n",
      "    %826 = mhlo.sine %824 : tensor<f32>\n",
      "    %827 = mhlo.negate %826 : tensor<f32>\n",
      "    %828 = mhlo.convert %825 : tensor<f32>\n",
      "    %829 = mhlo.convert %827 : tensor<f32>\n",
      "    %830 = \"mhlo.broadcast_in_dim\"(%828) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %831 = \"mhlo.broadcast_in_dim\"(%829) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %832 = \"mhlo.concatenate\"(%830, %831) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %833 = mhlo.convert %826 : tensor<f32>\n",
      "    %834 = mhlo.convert %825 : tensor<f32>\n",
      "    %835 = \"mhlo.broadcast_in_dim\"(%833) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %836 = \"mhlo.broadcast_in_dim\"(%834) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %837 = \"mhlo.concatenate\"(%835, %836) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %838 = \"mhlo.broadcast_in_dim\"(%832) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %839 = \"mhlo.broadcast_in_dim\"(%837) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %840 = \"mhlo.concatenate\"(%838, %839) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %841 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %842 = mhlo.constant dense<3.750000e-01> : tensor<f32>\n",
      "    %843 = mhlo.convert %841 : tensor<f32>\n",
      "    %844 = mhlo.convert %842 : tensor<f32>\n",
      "    %845 = mhlo.power %843, %844 : tensor<f32>\n",
      "    %846 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %847 = mhlo.divide %846, %845 : tensor<f32>\n",
      "    %848 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %849 = mhlo.multiply %848, %847 : tensor<f32>\n",
      "    %850 = mhlo.cosine %849 : tensor<f32>\n",
      "    %851 = mhlo.sine %849 : tensor<f32>\n",
      "    %852 = mhlo.negate %851 : tensor<f32>\n",
      "    %853 = mhlo.convert %850 : tensor<f32>\n",
      "    %854 = mhlo.convert %852 : tensor<f32>\n",
      "    %855 = \"mhlo.broadcast_in_dim\"(%853) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %856 = \"mhlo.broadcast_in_dim\"(%854) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %857 = \"mhlo.concatenate\"(%855, %856) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %858 = mhlo.convert %851 : tensor<f32>\n",
      "    %859 = mhlo.convert %850 : tensor<f32>\n",
      "    %860 = \"mhlo.broadcast_in_dim\"(%858) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %861 = \"mhlo.broadcast_in_dim\"(%859) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %862 = \"mhlo.concatenate\"(%860, %861) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %863 = \"mhlo.broadcast_in_dim\"(%857) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %864 = \"mhlo.broadcast_in_dim\"(%862) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %865 = \"mhlo.concatenate\"(%863, %864) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %866 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %867 = mhlo.constant dense<0.416666657> : tensor<f32>\n",
      "    %868 = mhlo.convert %866 : tensor<f32>\n",
      "    %869 = mhlo.convert %867 : tensor<f32>\n",
      "    %870 = mhlo.power %868, %869 : tensor<f32>\n",
      "    %871 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %872 = mhlo.divide %871, %870 : tensor<f32>\n",
      "    %873 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %874 = mhlo.multiply %873, %872 : tensor<f32>\n",
      "    %875 = mhlo.cosine %874 : tensor<f32>\n",
      "    %876 = mhlo.sine %874 : tensor<f32>\n",
      "    %877 = mhlo.negate %876 : tensor<f32>\n",
      "    %878 = mhlo.convert %875 : tensor<f32>\n",
      "    %879 = mhlo.convert %877 : tensor<f32>\n",
      "    %880 = \"mhlo.broadcast_in_dim\"(%878) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %881 = \"mhlo.broadcast_in_dim\"(%879) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %882 = \"mhlo.concatenate\"(%880, %881) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %883 = mhlo.convert %876 : tensor<f32>\n",
      "    %884 = mhlo.convert %875 : tensor<f32>\n",
      "    %885 = \"mhlo.broadcast_in_dim\"(%883) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %886 = \"mhlo.broadcast_in_dim\"(%884) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %887 = \"mhlo.concatenate\"(%885, %886) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %888 = \"mhlo.broadcast_in_dim\"(%882) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %889 = \"mhlo.broadcast_in_dim\"(%887) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %890 = \"mhlo.concatenate\"(%888, %889) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %891 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %892 = mhlo.constant dense<0.458333343> : tensor<f32>\n",
      "    %893 = mhlo.convert %891 : tensor<f32>\n",
      "    %894 = mhlo.convert %892 : tensor<f32>\n",
      "    %895 = mhlo.power %893, %894 : tensor<f32>\n",
      "    %896 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %897 = mhlo.divide %896, %895 : tensor<f32>\n",
      "    %898 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %899 = mhlo.multiply %898, %897 : tensor<f32>\n",
      "    %900 = mhlo.cosine %899 : tensor<f32>\n",
      "    %901 = mhlo.sine %899 : tensor<f32>\n",
      "    %902 = mhlo.negate %901 : tensor<f32>\n",
      "    %903 = mhlo.convert %900 : tensor<f32>\n",
      "    %904 = mhlo.convert %902 : tensor<f32>\n",
      "    %905 = \"mhlo.broadcast_in_dim\"(%903) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %906 = \"mhlo.broadcast_in_dim\"(%904) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %907 = \"mhlo.concatenate\"(%905, %906) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %908 = mhlo.convert %901 : tensor<f32>\n",
      "    %909 = mhlo.convert %900 : tensor<f32>\n",
      "    %910 = \"mhlo.broadcast_in_dim\"(%908) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %911 = \"mhlo.broadcast_in_dim\"(%909) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %912 = \"mhlo.concatenate\"(%910, %911) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %913 = \"mhlo.broadcast_in_dim\"(%907) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %914 = \"mhlo.broadcast_in_dim\"(%912) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %915 = \"mhlo.concatenate\"(%913, %914) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %916 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %917 = mhlo.constant dense<5.000000e-01> : tensor<f32>\n",
      "    %918 = mhlo.convert %916 : tensor<f32>\n",
      "    %919 = mhlo.convert %917 : tensor<f32>\n",
      "    %920 = mhlo.power %918, %919 : tensor<f32>\n",
      "    %921 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %922 = mhlo.divide %921, %920 : tensor<f32>\n",
      "    %923 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %924 = mhlo.multiply %923, %922 : tensor<f32>\n",
      "    %925 = mhlo.cosine %924 : tensor<f32>\n",
      "    %926 = mhlo.sine %924 : tensor<f32>\n",
      "    %927 = mhlo.negate %926 : tensor<f32>\n",
      "    %928 = mhlo.convert %925 : tensor<f32>\n",
      "    %929 = mhlo.convert %927 : tensor<f32>\n",
      "    %930 = \"mhlo.broadcast_in_dim\"(%928) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %931 = \"mhlo.broadcast_in_dim\"(%929) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %932 = \"mhlo.concatenate\"(%930, %931) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %933 = mhlo.convert %926 : tensor<f32>\n",
      "    %934 = mhlo.convert %925 : tensor<f32>\n",
      "    %935 = \"mhlo.broadcast_in_dim\"(%933) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %936 = \"mhlo.broadcast_in_dim\"(%934) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %937 = \"mhlo.concatenate\"(%935, %936) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %938 = \"mhlo.broadcast_in_dim\"(%932) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %939 = \"mhlo.broadcast_in_dim\"(%937) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %940 = \"mhlo.concatenate\"(%938, %939) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %941 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %942 = mhlo.constant dense<0.541666687> : tensor<f32>\n",
      "    %943 = mhlo.convert %941 : tensor<f32>\n",
      "    %944 = mhlo.convert %942 : tensor<f32>\n",
      "    %945 = mhlo.power %943, %944 : tensor<f32>\n",
      "    %946 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %947 = mhlo.divide %946, %945 : tensor<f32>\n",
      "    %948 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %949 = mhlo.multiply %948, %947 : tensor<f32>\n",
      "    %950 = mhlo.cosine %949 : tensor<f32>\n",
      "    %951 = mhlo.sine %949 : tensor<f32>\n",
      "    %952 = mhlo.negate %951 : tensor<f32>\n",
      "    %953 = mhlo.convert %950 : tensor<f32>\n",
      "    %954 = mhlo.convert %952 : tensor<f32>\n",
      "    %955 = \"mhlo.broadcast_in_dim\"(%953) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %956 = \"mhlo.broadcast_in_dim\"(%954) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %957 = \"mhlo.concatenate\"(%955, %956) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %958 = mhlo.convert %951 : tensor<f32>\n",
      "    %959 = mhlo.convert %950 : tensor<f32>\n",
      "    %960 = \"mhlo.broadcast_in_dim\"(%958) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %961 = \"mhlo.broadcast_in_dim\"(%959) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %962 = \"mhlo.concatenate\"(%960, %961) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %963 = \"mhlo.broadcast_in_dim\"(%957) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %964 = \"mhlo.broadcast_in_dim\"(%962) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %965 = \"mhlo.concatenate\"(%963, %964) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %966 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %967 = mhlo.constant dense<0.583333313> : tensor<f32>\n",
      "    %968 = mhlo.convert %966 : tensor<f32>\n",
      "    %969 = mhlo.convert %967 : tensor<f32>\n",
      "    %970 = mhlo.power %968, %969 : tensor<f32>\n",
      "    %971 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %972 = mhlo.divide %971, %970 : tensor<f32>\n",
      "    %973 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %974 = mhlo.multiply %973, %972 : tensor<f32>\n",
      "    %975 = mhlo.cosine %974 : tensor<f32>\n",
      "    %976 = mhlo.sine %974 : tensor<f32>\n",
      "    %977 = mhlo.negate %976 : tensor<f32>\n",
      "    %978 = mhlo.convert %975 : tensor<f32>\n",
      "    %979 = mhlo.convert %977 : tensor<f32>\n",
      "    %980 = \"mhlo.broadcast_in_dim\"(%978) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %981 = \"mhlo.broadcast_in_dim\"(%979) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %982 = \"mhlo.concatenate\"(%980, %981) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %983 = mhlo.convert %976 : tensor<f32>\n",
      "    %984 = mhlo.convert %975 : tensor<f32>\n",
      "    %985 = \"mhlo.broadcast_in_dim\"(%983) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %986 = \"mhlo.broadcast_in_dim\"(%984) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %987 = \"mhlo.concatenate\"(%985, %986) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %988 = \"mhlo.broadcast_in_dim\"(%982) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %989 = \"mhlo.broadcast_in_dim\"(%987) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %990 = \"mhlo.concatenate\"(%988, %989) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %991 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %992 = mhlo.constant dense<6.250000e-01> : tensor<f32>\n",
      "    %993 = mhlo.convert %991 : tensor<f32>\n",
      "    %994 = mhlo.convert %992 : tensor<f32>\n",
      "    %995 = mhlo.power %993, %994 : tensor<f32>\n",
      "    %996 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %997 = mhlo.divide %996, %995 : tensor<f32>\n",
      "    %998 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %999 = mhlo.multiply %998, %997 : tensor<f32>\n",
      "    %1000 = mhlo.cosine %999 : tensor<f32>\n",
      "    %1001 = mhlo.sine %999 : tensor<f32>\n",
      "    %1002 = mhlo.negate %1001 : tensor<f32>\n",
      "    %1003 = mhlo.convert %1000 : tensor<f32>\n",
      "    %1004 = mhlo.convert %1002 : tensor<f32>\n",
      "    %1005 = \"mhlo.broadcast_in_dim\"(%1003) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1006 = \"mhlo.broadcast_in_dim\"(%1004) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1007 = \"mhlo.concatenate\"(%1005, %1006) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1008 = mhlo.convert %1001 : tensor<f32>\n",
      "    %1009 = mhlo.convert %1000 : tensor<f32>\n",
      "    %1010 = \"mhlo.broadcast_in_dim\"(%1008) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1011 = \"mhlo.broadcast_in_dim\"(%1009) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1012 = \"mhlo.concatenate\"(%1010, %1011) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1013 = \"mhlo.broadcast_in_dim\"(%1007) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1014 = \"mhlo.broadcast_in_dim\"(%1012) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1015 = \"mhlo.concatenate\"(%1013, %1014) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1016 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1017 = mhlo.constant dense<0.666666686> : tensor<f32>\n",
      "    %1018 = mhlo.convert %1016 : tensor<f32>\n",
      "    %1019 = mhlo.convert %1017 : tensor<f32>\n",
      "    %1020 = mhlo.power %1018, %1019 : tensor<f32>\n",
      "    %1021 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1022 = mhlo.divide %1021, %1020 : tensor<f32>\n",
      "    %1023 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1024 = mhlo.multiply %1023, %1022 : tensor<f32>\n",
      "    %1025 = mhlo.cosine %1024 : tensor<f32>\n",
      "    %1026 = mhlo.sine %1024 : tensor<f32>\n",
      "    %1027 = mhlo.negate %1026 : tensor<f32>\n",
      "    %1028 = mhlo.convert %1025 : tensor<f32>\n",
      "    %1029 = mhlo.convert %1027 : tensor<f32>\n",
      "    %1030 = \"mhlo.broadcast_in_dim\"(%1028) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1031 = \"mhlo.broadcast_in_dim\"(%1029) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1032 = \"mhlo.concatenate\"(%1030, %1031) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1033 = mhlo.convert %1026 : tensor<f32>\n",
      "    %1034 = mhlo.convert %1025 : tensor<f32>\n",
      "    %1035 = \"mhlo.broadcast_in_dim\"(%1033) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1036 = \"mhlo.broadcast_in_dim\"(%1034) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1037 = \"mhlo.concatenate\"(%1035, %1036) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1038 = \"mhlo.broadcast_in_dim\"(%1032) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1039 = \"mhlo.broadcast_in_dim\"(%1037) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1040 = \"mhlo.concatenate\"(%1038, %1039) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1041 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1042 = mhlo.constant dense<0.708333313> : tensor<f32>\n",
      "    %1043 = mhlo.convert %1041 : tensor<f32>\n",
      "    %1044 = mhlo.convert %1042 : tensor<f32>\n",
      "    %1045 = mhlo.power %1043, %1044 : tensor<f32>\n",
      "    %1046 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1047 = mhlo.divide %1046, %1045 : tensor<f32>\n",
      "    %1048 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1049 = mhlo.multiply %1048, %1047 : tensor<f32>\n",
      "    %1050 = mhlo.cosine %1049 : tensor<f32>\n",
      "    %1051 = mhlo.sine %1049 : tensor<f32>\n",
      "    %1052 = mhlo.negate %1051 : tensor<f32>\n",
      "    %1053 = mhlo.convert %1050 : tensor<f32>\n",
      "    %1054 = mhlo.convert %1052 : tensor<f32>\n",
      "    %1055 = \"mhlo.broadcast_in_dim\"(%1053) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1056 = \"mhlo.broadcast_in_dim\"(%1054) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1057 = \"mhlo.concatenate\"(%1055, %1056) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1058 = mhlo.convert %1051 : tensor<f32>\n",
      "    %1059 = mhlo.convert %1050 : tensor<f32>\n",
      "    %1060 = \"mhlo.broadcast_in_dim\"(%1058) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1061 = \"mhlo.broadcast_in_dim\"(%1059) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1062 = \"mhlo.concatenate\"(%1060, %1061) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1063 = \"mhlo.broadcast_in_dim\"(%1057) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1064 = \"mhlo.broadcast_in_dim\"(%1062) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1065 = \"mhlo.concatenate\"(%1063, %1064) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1066 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1067 = mhlo.constant dense<7.500000e-01> : tensor<f32>\n",
      "    %1068 = mhlo.convert %1066 : tensor<f32>\n",
      "    %1069 = mhlo.convert %1067 : tensor<f32>\n",
      "    %1070 = mhlo.power %1068, %1069 : tensor<f32>\n",
      "    %1071 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1072 = mhlo.divide %1071, %1070 : tensor<f32>\n",
      "    %1073 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1074 = mhlo.multiply %1073, %1072 : tensor<f32>\n",
      "    %1075 = mhlo.cosine %1074 : tensor<f32>\n",
      "    %1076 = mhlo.sine %1074 : tensor<f32>\n",
      "    %1077 = mhlo.negate %1076 : tensor<f32>\n",
      "    %1078 = mhlo.convert %1075 : tensor<f32>\n",
      "    %1079 = mhlo.convert %1077 : tensor<f32>\n",
      "    %1080 = \"mhlo.broadcast_in_dim\"(%1078) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1081 = \"mhlo.broadcast_in_dim\"(%1079) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1082 = \"mhlo.concatenate\"(%1080, %1081) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1083 = mhlo.convert %1076 : tensor<f32>\n",
      "    %1084 = mhlo.convert %1075 : tensor<f32>\n",
      "    %1085 = \"mhlo.broadcast_in_dim\"(%1083) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1086 = \"mhlo.broadcast_in_dim\"(%1084) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1087 = \"mhlo.concatenate\"(%1085, %1086) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1088 = \"mhlo.broadcast_in_dim\"(%1082) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1089 = \"mhlo.broadcast_in_dim\"(%1087) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1090 = \"mhlo.concatenate\"(%1088, %1089) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1091 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1092 = mhlo.constant dense<0.791666686> : tensor<f32>\n",
      "    %1093 = mhlo.convert %1091 : tensor<f32>\n",
      "    %1094 = mhlo.convert %1092 : tensor<f32>\n",
      "    %1095 = mhlo.power %1093, %1094 : tensor<f32>\n",
      "    %1096 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1097 = mhlo.divide %1096, %1095 : tensor<f32>\n",
      "    %1098 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1099 = mhlo.multiply %1098, %1097 : tensor<f32>\n",
      "    %1100 = mhlo.cosine %1099 : tensor<f32>\n",
      "    %1101 = mhlo.sine %1099 : tensor<f32>\n",
      "    %1102 = mhlo.negate %1101 : tensor<f32>\n",
      "    %1103 = mhlo.convert %1100 : tensor<f32>\n",
      "    %1104 = mhlo.convert %1102 : tensor<f32>\n",
      "    %1105 = \"mhlo.broadcast_in_dim\"(%1103) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1106 = \"mhlo.broadcast_in_dim\"(%1104) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1107 = \"mhlo.concatenate\"(%1105, %1106) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1108 = mhlo.convert %1101 : tensor<f32>\n",
      "    %1109 = mhlo.convert %1100 : tensor<f32>\n",
      "    %1110 = \"mhlo.broadcast_in_dim\"(%1108) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1111 = \"mhlo.broadcast_in_dim\"(%1109) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1112 = \"mhlo.concatenate\"(%1110, %1111) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1113 = \"mhlo.broadcast_in_dim\"(%1107) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1114 = \"mhlo.broadcast_in_dim\"(%1112) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1115 = \"mhlo.concatenate\"(%1113, %1114) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1116 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1117 = mhlo.constant dense<0.833333313> : tensor<f32>\n",
      "    %1118 = mhlo.convert %1116 : tensor<f32>\n",
      "    %1119 = mhlo.convert %1117 : tensor<f32>\n",
      "    %1120 = mhlo.power %1118, %1119 : tensor<f32>\n",
      "    %1121 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1122 = mhlo.divide %1121, %1120 : tensor<f32>\n",
      "    %1123 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1124 = mhlo.multiply %1123, %1122 : tensor<f32>\n",
      "    %1125 = mhlo.cosine %1124 : tensor<f32>\n",
      "    %1126 = mhlo.sine %1124 : tensor<f32>\n",
      "    %1127 = mhlo.negate %1126 : tensor<f32>\n",
      "    %1128 = mhlo.convert %1125 : tensor<f32>\n",
      "    %1129 = mhlo.convert %1127 : tensor<f32>\n",
      "    %1130 = \"mhlo.broadcast_in_dim\"(%1128) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1131 = \"mhlo.broadcast_in_dim\"(%1129) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1132 = \"mhlo.concatenate\"(%1130, %1131) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1133 = mhlo.convert %1126 : tensor<f32>\n",
      "    %1134 = mhlo.convert %1125 : tensor<f32>\n",
      "    %1135 = \"mhlo.broadcast_in_dim\"(%1133) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1136 = \"mhlo.broadcast_in_dim\"(%1134) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1137 = \"mhlo.concatenate\"(%1135, %1136) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1138 = \"mhlo.broadcast_in_dim\"(%1132) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1139 = \"mhlo.broadcast_in_dim\"(%1137) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1140 = \"mhlo.concatenate\"(%1138, %1139) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1141 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1142 = mhlo.constant dense<8.750000e-01> : tensor<f32>\n",
      "    %1143 = mhlo.convert %1141 : tensor<f32>\n",
      "    %1144 = mhlo.convert %1142 : tensor<f32>\n",
      "    %1145 = mhlo.power %1143, %1144 : tensor<f32>\n",
      "    %1146 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1147 = mhlo.divide %1146, %1145 : tensor<f32>\n",
      "    %1148 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1149 = mhlo.multiply %1148, %1147 : tensor<f32>\n",
      "    %1150 = mhlo.cosine %1149 : tensor<f32>\n",
      "    %1151 = mhlo.sine %1149 : tensor<f32>\n",
      "    %1152 = mhlo.negate %1151 : tensor<f32>\n",
      "    %1153 = mhlo.convert %1150 : tensor<f32>\n",
      "    %1154 = mhlo.convert %1152 : tensor<f32>\n",
      "    %1155 = \"mhlo.broadcast_in_dim\"(%1153) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1156 = \"mhlo.broadcast_in_dim\"(%1154) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1157 = \"mhlo.concatenate\"(%1155, %1156) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1158 = mhlo.convert %1151 : tensor<f32>\n",
      "    %1159 = mhlo.convert %1150 : tensor<f32>\n",
      "    %1160 = \"mhlo.broadcast_in_dim\"(%1158) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1161 = \"mhlo.broadcast_in_dim\"(%1159) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1162 = \"mhlo.concatenate\"(%1160, %1161) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1163 = \"mhlo.broadcast_in_dim\"(%1157) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1164 = \"mhlo.broadcast_in_dim\"(%1162) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1165 = \"mhlo.concatenate\"(%1163, %1164) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1166 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1167 = mhlo.constant dense<0.916666686> : tensor<f32>\n",
      "    %1168 = mhlo.convert %1166 : tensor<f32>\n",
      "    %1169 = mhlo.convert %1167 : tensor<f32>\n",
      "    %1170 = mhlo.power %1168, %1169 : tensor<f32>\n",
      "    %1171 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1172 = mhlo.divide %1171, %1170 : tensor<f32>\n",
      "    %1173 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1174 = mhlo.multiply %1173, %1172 : tensor<f32>\n",
      "    %1175 = mhlo.cosine %1174 : tensor<f32>\n",
      "    %1176 = mhlo.sine %1174 : tensor<f32>\n",
      "    %1177 = mhlo.negate %1176 : tensor<f32>\n",
      "    %1178 = mhlo.convert %1175 : tensor<f32>\n",
      "    %1179 = mhlo.convert %1177 : tensor<f32>\n",
      "    %1180 = \"mhlo.broadcast_in_dim\"(%1178) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1181 = \"mhlo.broadcast_in_dim\"(%1179) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1182 = \"mhlo.concatenate\"(%1180, %1181) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1183 = mhlo.convert %1176 : tensor<f32>\n",
      "    %1184 = mhlo.convert %1175 : tensor<f32>\n",
      "    %1185 = \"mhlo.broadcast_in_dim\"(%1183) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1186 = \"mhlo.broadcast_in_dim\"(%1184) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1187 = \"mhlo.concatenate\"(%1185, %1186) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1188 = \"mhlo.broadcast_in_dim\"(%1182) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1189 = \"mhlo.broadcast_in_dim\"(%1187) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1190 = \"mhlo.concatenate\"(%1188, %1189) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1191 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1192 = mhlo.constant dense<0.958333313> : tensor<f32>\n",
      "    %1193 = mhlo.convert %1191 : tensor<f32>\n",
      "    %1194 = mhlo.convert %1192 : tensor<f32>\n",
      "    %1195 = mhlo.power %1193, %1194 : tensor<f32>\n",
      "    %1196 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1197 = mhlo.divide %1196, %1195 : tensor<f32>\n",
      "    %1198 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1199 = mhlo.multiply %1198, %1197 : tensor<f32>\n",
      "    %1200 = mhlo.cosine %1199 : tensor<f32>\n",
      "    %1201 = mhlo.sine %1199 : tensor<f32>\n",
      "    %1202 = mhlo.negate %1201 : tensor<f32>\n",
      "    %1203 = mhlo.convert %1200 : tensor<f32>\n",
      "    %1204 = mhlo.convert %1202 : tensor<f32>\n",
      "    %1205 = \"mhlo.broadcast_in_dim\"(%1203) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1206 = \"mhlo.broadcast_in_dim\"(%1204) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1207 = \"mhlo.concatenate\"(%1205, %1206) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1208 = mhlo.convert %1201 : tensor<f32>\n",
      "    %1209 = mhlo.convert %1200 : tensor<f32>\n",
      "    %1210 = \"mhlo.broadcast_in_dim\"(%1208) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1211 = \"mhlo.broadcast_in_dim\"(%1209) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1212 = \"mhlo.concatenate\"(%1210, %1211) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1213 = \"mhlo.broadcast_in_dim\"(%1207) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1214 = \"mhlo.broadcast_in_dim\"(%1212) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1215 = \"mhlo.concatenate\"(%1213, %1214) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1216 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1217 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1218 = mhlo.convert %1216 : tensor<f32>\n",
      "    %1219 = mhlo.convert %1217 : tensor<f32>\n",
      "    %1220 = mhlo.power %1218, %1219 : tensor<f32>\n",
      "    %1221 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1222 = mhlo.divide %1221, %1220 : tensor<f32>\n",
      "    %1223 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1224 = mhlo.multiply %1223, %1222 : tensor<f32>\n",
      "    %1225 = mhlo.cosine %1224 : tensor<f32>\n",
      "    %1226 = mhlo.sine %1224 : tensor<f32>\n",
      "    %1227 = mhlo.negate %1226 : tensor<f32>\n",
      "    %1228 = mhlo.convert %1225 : tensor<f32>\n",
      "    %1229 = mhlo.convert %1227 : tensor<f32>\n",
      "    %1230 = \"mhlo.broadcast_in_dim\"(%1228) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1231 = \"mhlo.broadcast_in_dim\"(%1229) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1232 = \"mhlo.concatenate\"(%1230, %1231) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1233 = mhlo.convert %1226 : tensor<f32>\n",
      "    %1234 = mhlo.convert %1225 : tensor<f32>\n",
      "    %1235 = \"mhlo.broadcast_in_dim\"(%1233) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1236 = \"mhlo.broadcast_in_dim\"(%1234) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1237 = \"mhlo.concatenate\"(%1235, %1236) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1238 = \"mhlo.broadcast_in_dim\"(%1232) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1239 = \"mhlo.broadcast_in_dim\"(%1237) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1240 = \"mhlo.concatenate\"(%1238, %1239) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1241 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1242 = mhlo.constant dense<0.0416666679> : tensor<f32>\n",
      "    %1243 = mhlo.convert %1241 : tensor<f32>\n",
      "    %1244 = mhlo.convert %1242 : tensor<f32>\n",
      "    %1245 = mhlo.power %1243, %1244 : tensor<f32>\n",
      "    %1246 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1247 = mhlo.divide %1246, %1245 : tensor<f32>\n",
      "    %1248 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1249 = mhlo.multiply %1248, %1247 : tensor<f32>\n",
      "    %1250 = mhlo.cosine %1249 : tensor<f32>\n",
      "    %1251 = mhlo.sine %1249 : tensor<f32>\n",
      "    %1252 = mhlo.negate %1251 : tensor<f32>\n",
      "    %1253 = mhlo.convert %1250 : tensor<f32>\n",
      "    %1254 = mhlo.convert %1252 : tensor<f32>\n",
      "    %1255 = \"mhlo.broadcast_in_dim\"(%1253) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1256 = \"mhlo.broadcast_in_dim\"(%1254) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1257 = \"mhlo.concatenate\"(%1255, %1256) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1258 = mhlo.convert %1251 : tensor<f32>\n",
      "    %1259 = mhlo.convert %1250 : tensor<f32>\n",
      "    %1260 = \"mhlo.broadcast_in_dim\"(%1258) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1261 = \"mhlo.broadcast_in_dim\"(%1259) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1262 = \"mhlo.concatenate\"(%1260, %1261) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1263 = \"mhlo.broadcast_in_dim\"(%1257) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1264 = \"mhlo.broadcast_in_dim\"(%1262) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1265 = \"mhlo.concatenate\"(%1263, %1264) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1266 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1267 = mhlo.constant dense<0.0833333358> : tensor<f32>\n",
      "    %1268 = mhlo.convert %1266 : tensor<f32>\n",
      "    %1269 = mhlo.convert %1267 : tensor<f32>\n",
      "    %1270 = mhlo.power %1268, %1269 : tensor<f32>\n",
      "    %1271 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1272 = mhlo.divide %1271, %1270 : tensor<f32>\n",
      "    %1273 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1274 = mhlo.multiply %1273, %1272 : tensor<f32>\n",
      "    %1275 = mhlo.cosine %1274 : tensor<f32>\n",
      "    %1276 = mhlo.sine %1274 : tensor<f32>\n",
      "    %1277 = mhlo.negate %1276 : tensor<f32>\n",
      "    %1278 = mhlo.convert %1275 : tensor<f32>\n",
      "    %1279 = mhlo.convert %1277 : tensor<f32>\n",
      "    %1280 = \"mhlo.broadcast_in_dim\"(%1278) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1281 = \"mhlo.broadcast_in_dim\"(%1279) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1282 = \"mhlo.concatenate\"(%1280, %1281) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1283 = mhlo.convert %1276 : tensor<f32>\n",
      "    %1284 = mhlo.convert %1275 : tensor<f32>\n",
      "    %1285 = \"mhlo.broadcast_in_dim\"(%1283) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1286 = \"mhlo.broadcast_in_dim\"(%1284) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1287 = \"mhlo.concatenate\"(%1285, %1286) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1288 = \"mhlo.broadcast_in_dim\"(%1282) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1289 = \"mhlo.broadcast_in_dim\"(%1287) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1290 = \"mhlo.concatenate\"(%1288, %1289) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1291 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1292 = mhlo.constant dense<1.250000e-01> : tensor<f32>\n",
      "    %1293 = mhlo.convert %1291 : tensor<f32>\n",
      "    %1294 = mhlo.convert %1292 : tensor<f32>\n",
      "    %1295 = mhlo.power %1293, %1294 : tensor<f32>\n",
      "    %1296 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1297 = mhlo.divide %1296, %1295 : tensor<f32>\n",
      "    %1298 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1299 = mhlo.multiply %1298, %1297 : tensor<f32>\n",
      "    %1300 = mhlo.cosine %1299 : tensor<f32>\n",
      "    %1301 = mhlo.sine %1299 : tensor<f32>\n",
      "    %1302 = mhlo.negate %1301 : tensor<f32>\n",
      "    %1303 = mhlo.convert %1300 : tensor<f32>\n",
      "    %1304 = mhlo.convert %1302 : tensor<f32>\n",
      "    %1305 = \"mhlo.broadcast_in_dim\"(%1303) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1306 = \"mhlo.broadcast_in_dim\"(%1304) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1307 = \"mhlo.concatenate\"(%1305, %1306) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1308 = mhlo.convert %1301 : tensor<f32>\n",
      "    %1309 = mhlo.convert %1300 : tensor<f32>\n",
      "    %1310 = \"mhlo.broadcast_in_dim\"(%1308) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1311 = \"mhlo.broadcast_in_dim\"(%1309) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1312 = \"mhlo.concatenate\"(%1310, %1311) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1313 = \"mhlo.broadcast_in_dim\"(%1307) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1314 = \"mhlo.broadcast_in_dim\"(%1312) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1315 = \"mhlo.concatenate\"(%1313, %1314) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1316 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1317 = mhlo.constant dense<0.166666672> : tensor<f32>\n",
      "    %1318 = mhlo.convert %1316 : tensor<f32>\n",
      "    %1319 = mhlo.convert %1317 : tensor<f32>\n",
      "    %1320 = mhlo.power %1318, %1319 : tensor<f32>\n",
      "    %1321 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1322 = mhlo.divide %1321, %1320 : tensor<f32>\n",
      "    %1323 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1324 = mhlo.multiply %1323, %1322 : tensor<f32>\n",
      "    %1325 = mhlo.cosine %1324 : tensor<f32>\n",
      "    %1326 = mhlo.sine %1324 : tensor<f32>\n",
      "    %1327 = mhlo.negate %1326 : tensor<f32>\n",
      "    %1328 = mhlo.convert %1325 : tensor<f32>\n",
      "    %1329 = mhlo.convert %1327 : tensor<f32>\n",
      "    %1330 = \"mhlo.broadcast_in_dim\"(%1328) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1331 = \"mhlo.broadcast_in_dim\"(%1329) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1332 = \"mhlo.concatenate\"(%1330, %1331) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1333 = mhlo.convert %1326 : tensor<f32>\n",
      "    %1334 = mhlo.convert %1325 : tensor<f32>\n",
      "    %1335 = \"mhlo.broadcast_in_dim\"(%1333) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1336 = \"mhlo.broadcast_in_dim\"(%1334) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1337 = \"mhlo.concatenate\"(%1335, %1336) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1338 = \"mhlo.broadcast_in_dim\"(%1332) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1339 = \"mhlo.broadcast_in_dim\"(%1337) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1340 = \"mhlo.concatenate\"(%1338, %1339) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1341 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1342 = mhlo.constant dense<0.208333328> : tensor<f32>\n",
      "    %1343 = mhlo.convert %1341 : tensor<f32>\n",
      "    %1344 = mhlo.convert %1342 : tensor<f32>\n",
      "    %1345 = mhlo.power %1343, %1344 : tensor<f32>\n",
      "    %1346 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1347 = mhlo.divide %1346, %1345 : tensor<f32>\n",
      "    %1348 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1349 = mhlo.multiply %1348, %1347 : tensor<f32>\n",
      "    %1350 = mhlo.cosine %1349 : tensor<f32>\n",
      "    %1351 = mhlo.sine %1349 : tensor<f32>\n",
      "    %1352 = mhlo.negate %1351 : tensor<f32>\n",
      "    %1353 = mhlo.convert %1350 : tensor<f32>\n",
      "    %1354 = mhlo.convert %1352 : tensor<f32>\n",
      "    %1355 = \"mhlo.broadcast_in_dim\"(%1353) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1356 = \"mhlo.broadcast_in_dim\"(%1354) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1357 = \"mhlo.concatenate\"(%1355, %1356) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1358 = mhlo.convert %1351 : tensor<f32>\n",
      "    %1359 = mhlo.convert %1350 : tensor<f32>\n",
      "    %1360 = \"mhlo.broadcast_in_dim\"(%1358) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1361 = \"mhlo.broadcast_in_dim\"(%1359) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1362 = \"mhlo.concatenate\"(%1360, %1361) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1363 = \"mhlo.broadcast_in_dim\"(%1357) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1364 = \"mhlo.broadcast_in_dim\"(%1362) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1365 = \"mhlo.concatenate\"(%1363, %1364) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1366 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1367 = mhlo.constant dense<2.500000e-01> : tensor<f32>\n",
      "    %1368 = mhlo.convert %1366 : tensor<f32>\n",
      "    %1369 = mhlo.convert %1367 : tensor<f32>\n",
      "    %1370 = mhlo.power %1368, %1369 : tensor<f32>\n",
      "    %1371 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1372 = mhlo.divide %1371, %1370 : tensor<f32>\n",
      "    %1373 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1374 = mhlo.multiply %1373, %1372 : tensor<f32>\n",
      "    %1375 = mhlo.cosine %1374 : tensor<f32>\n",
      "    %1376 = mhlo.sine %1374 : tensor<f32>\n",
      "    %1377 = mhlo.negate %1376 : tensor<f32>\n",
      "    %1378 = mhlo.convert %1375 : tensor<f32>\n",
      "    %1379 = mhlo.convert %1377 : tensor<f32>\n",
      "    %1380 = \"mhlo.broadcast_in_dim\"(%1378) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1381 = \"mhlo.broadcast_in_dim\"(%1379) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1382 = \"mhlo.concatenate\"(%1380, %1381) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1383 = mhlo.convert %1376 : tensor<f32>\n",
      "    %1384 = mhlo.convert %1375 : tensor<f32>\n",
      "    %1385 = \"mhlo.broadcast_in_dim\"(%1383) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1386 = \"mhlo.broadcast_in_dim\"(%1384) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1387 = \"mhlo.concatenate\"(%1385, %1386) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1388 = \"mhlo.broadcast_in_dim\"(%1382) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1389 = \"mhlo.broadcast_in_dim\"(%1387) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1390 = \"mhlo.concatenate\"(%1388, %1389) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1391 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1392 = mhlo.constant dense<0.291666657> : tensor<f32>\n",
      "    %1393 = mhlo.convert %1391 : tensor<f32>\n",
      "    %1394 = mhlo.convert %1392 : tensor<f32>\n",
      "    %1395 = mhlo.power %1393, %1394 : tensor<f32>\n",
      "    %1396 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1397 = mhlo.divide %1396, %1395 : tensor<f32>\n",
      "    %1398 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1399 = mhlo.multiply %1398, %1397 : tensor<f32>\n",
      "    %1400 = mhlo.cosine %1399 : tensor<f32>\n",
      "    %1401 = mhlo.sine %1399 : tensor<f32>\n",
      "    %1402 = mhlo.negate %1401 : tensor<f32>\n",
      "    %1403 = mhlo.convert %1400 : tensor<f32>\n",
      "    %1404 = mhlo.convert %1402 : tensor<f32>\n",
      "    %1405 = \"mhlo.broadcast_in_dim\"(%1403) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1406 = \"mhlo.broadcast_in_dim\"(%1404) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1407 = \"mhlo.concatenate\"(%1405, %1406) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1408 = mhlo.convert %1401 : tensor<f32>\n",
      "    %1409 = mhlo.convert %1400 : tensor<f32>\n",
      "    %1410 = \"mhlo.broadcast_in_dim\"(%1408) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1411 = \"mhlo.broadcast_in_dim\"(%1409) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1412 = \"mhlo.concatenate\"(%1410, %1411) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1413 = \"mhlo.broadcast_in_dim\"(%1407) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1414 = \"mhlo.broadcast_in_dim\"(%1412) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1415 = \"mhlo.concatenate\"(%1413, %1414) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1416 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1417 = mhlo.constant dense<0.333333343> : tensor<f32>\n",
      "    %1418 = mhlo.convert %1416 : tensor<f32>\n",
      "    %1419 = mhlo.convert %1417 : tensor<f32>\n",
      "    %1420 = mhlo.power %1418, %1419 : tensor<f32>\n",
      "    %1421 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1422 = mhlo.divide %1421, %1420 : tensor<f32>\n",
      "    %1423 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1424 = mhlo.multiply %1423, %1422 : tensor<f32>\n",
      "    %1425 = mhlo.cosine %1424 : tensor<f32>\n",
      "    %1426 = mhlo.sine %1424 : tensor<f32>\n",
      "    %1427 = mhlo.negate %1426 : tensor<f32>\n",
      "    %1428 = mhlo.convert %1425 : tensor<f32>\n",
      "    %1429 = mhlo.convert %1427 : tensor<f32>\n",
      "    %1430 = \"mhlo.broadcast_in_dim\"(%1428) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1431 = \"mhlo.broadcast_in_dim\"(%1429) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1432 = \"mhlo.concatenate\"(%1430, %1431) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1433 = mhlo.convert %1426 : tensor<f32>\n",
      "    %1434 = mhlo.convert %1425 : tensor<f32>\n",
      "    %1435 = \"mhlo.broadcast_in_dim\"(%1433) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1436 = \"mhlo.broadcast_in_dim\"(%1434) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1437 = \"mhlo.concatenate\"(%1435, %1436) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1438 = \"mhlo.broadcast_in_dim\"(%1432) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1439 = \"mhlo.broadcast_in_dim\"(%1437) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1440 = \"mhlo.concatenate\"(%1438, %1439) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1441 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1442 = mhlo.constant dense<3.750000e-01> : tensor<f32>\n",
      "    %1443 = mhlo.convert %1441 : tensor<f32>\n",
      "    %1444 = mhlo.convert %1442 : tensor<f32>\n",
      "    %1445 = mhlo.power %1443, %1444 : tensor<f32>\n",
      "    %1446 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1447 = mhlo.divide %1446, %1445 : tensor<f32>\n",
      "    %1448 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1449 = mhlo.multiply %1448, %1447 : tensor<f32>\n",
      "    %1450 = mhlo.cosine %1449 : tensor<f32>\n",
      "    %1451 = mhlo.sine %1449 : tensor<f32>\n",
      "    %1452 = mhlo.negate %1451 : tensor<f32>\n",
      "    %1453 = mhlo.convert %1450 : tensor<f32>\n",
      "    %1454 = mhlo.convert %1452 : tensor<f32>\n",
      "    %1455 = \"mhlo.broadcast_in_dim\"(%1453) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1456 = \"mhlo.broadcast_in_dim\"(%1454) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1457 = \"mhlo.concatenate\"(%1455, %1456) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1458 = mhlo.convert %1451 : tensor<f32>\n",
      "    %1459 = mhlo.convert %1450 : tensor<f32>\n",
      "    %1460 = \"mhlo.broadcast_in_dim\"(%1458) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1461 = \"mhlo.broadcast_in_dim\"(%1459) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1462 = \"mhlo.concatenate\"(%1460, %1461) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1463 = \"mhlo.broadcast_in_dim\"(%1457) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1464 = \"mhlo.broadcast_in_dim\"(%1462) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1465 = \"mhlo.concatenate\"(%1463, %1464) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1466 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1467 = mhlo.constant dense<0.416666657> : tensor<f32>\n",
      "    %1468 = mhlo.convert %1466 : tensor<f32>\n",
      "    %1469 = mhlo.convert %1467 : tensor<f32>\n",
      "    %1470 = mhlo.power %1468, %1469 : tensor<f32>\n",
      "    %1471 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1472 = mhlo.divide %1471, %1470 : tensor<f32>\n",
      "    %1473 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1474 = mhlo.multiply %1473, %1472 : tensor<f32>\n",
      "    %1475 = mhlo.cosine %1474 : tensor<f32>\n",
      "    %1476 = mhlo.sine %1474 : tensor<f32>\n",
      "    %1477 = mhlo.negate %1476 : tensor<f32>\n",
      "    %1478 = mhlo.convert %1475 : tensor<f32>\n",
      "    %1479 = mhlo.convert %1477 : tensor<f32>\n",
      "    %1480 = \"mhlo.broadcast_in_dim\"(%1478) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1481 = \"mhlo.broadcast_in_dim\"(%1479) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1482 = \"mhlo.concatenate\"(%1480, %1481) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1483 = mhlo.convert %1476 : tensor<f32>\n",
      "    %1484 = mhlo.convert %1475 : tensor<f32>\n",
      "    %1485 = \"mhlo.broadcast_in_dim\"(%1483) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1486 = \"mhlo.broadcast_in_dim\"(%1484) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1487 = \"mhlo.concatenate\"(%1485, %1486) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1488 = \"mhlo.broadcast_in_dim\"(%1482) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1489 = \"mhlo.broadcast_in_dim\"(%1487) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1490 = \"mhlo.concatenate\"(%1488, %1489) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1491 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1492 = mhlo.constant dense<0.458333343> : tensor<f32>\n",
      "    %1493 = mhlo.convert %1491 : tensor<f32>\n",
      "    %1494 = mhlo.convert %1492 : tensor<f32>\n",
      "    %1495 = mhlo.power %1493, %1494 : tensor<f32>\n",
      "    %1496 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1497 = mhlo.divide %1496, %1495 : tensor<f32>\n",
      "    %1498 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1499 = mhlo.multiply %1498, %1497 : tensor<f32>\n",
      "    %1500 = mhlo.cosine %1499 : tensor<f32>\n",
      "    %1501 = mhlo.sine %1499 : tensor<f32>\n",
      "    %1502 = mhlo.negate %1501 : tensor<f32>\n",
      "    %1503 = mhlo.convert %1500 : tensor<f32>\n",
      "    %1504 = mhlo.convert %1502 : tensor<f32>\n",
      "    %1505 = \"mhlo.broadcast_in_dim\"(%1503) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1506 = \"mhlo.broadcast_in_dim\"(%1504) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1507 = \"mhlo.concatenate\"(%1505, %1506) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1508 = mhlo.convert %1501 : tensor<f32>\n",
      "    %1509 = mhlo.convert %1500 : tensor<f32>\n",
      "    %1510 = \"mhlo.broadcast_in_dim\"(%1508) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1511 = \"mhlo.broadcast_in_dim\"(%1509) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1512 = \"mhlo.concatenate\"(%1510, %1511) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1513 = \"mhlo.broadcast_in_dim\"(%1507) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1514 = \"mhlo.broadcast_in_dim\"(%1512) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1515 = \"mhlo.concatenate\"(%1513, %1514) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1516 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1517 = mhlo.constant dense<5.000000e-01> : tensor<f32>\n",
      "    %1518 = mhlo.convert %1516 : tensor<f32>\n",
      "    %1519 = mhlo.convert %1517 : tensor<f32>\n",
      "    %1520 = mhlo.power %1518, %1519 : tensor<f32>\n",
      "    %1521 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1522 = mhlo.divide %1521, %1520 : tensor<f32>\n",
      "    %1523 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1524 = mhlo.multiply %1523, %1522 : tensor<f32>\n",
      "    %1525 = mhlo.cosine %1524 : tensor<f32>\n",
      "    %1526 = mhlo.sine %1524 : tensor<f32>\n",
      "    %1527 = mhlo.negate %1526 : tensor<f32>\n",
      "    %1528 = mhlo.convert %1525 : tensor<f32>\n",
      "    %1529 = mhlo.convert %1527 : tensor<f32>\n",
      "    %1530 = \"mhlo.broadcast_in_dim\"(%1528) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1531 = \"mhlo.broadcast_in_dim\"(%1529) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1532 = \"mhlo.concatenate\"(%1530, %1531) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1533 = mhlo.convert %1526 : tensor<f32>\n",
      "    %1534 = mhlo.convert %1525 : tensor<f32>\n",
      "    %1535 = \"mhlo.broadcast_in_dim\"(%1533) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1536 = \"mhlo.broadcast_in_dim\"(%1534) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1537 = \"mhlo.concatenate\"(%1535, %1536) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1538 = \"mhlo.broadcast_in_dim\"(%1532) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1539 = \"mhlo.broadcast_in_dim\"(%1537) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1540 = \"mhlo.concatenate\"(%1538, %1539) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1541 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1542 = mhlo.constant dense<0.541666687> : tensor<f32>\n",
      "    %1543 = mhlo.convert %1541 : tensor<f32>\n",
      "    %1544 = mhlo.convert %1542 : tensor<f32>\n",
      "    %1545 = mhlo.power %1543, %1544 : tensor<f32>\n",
      "    %1546 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1547 = mhlo.divide %1546, %1545 : tensor<f32>\n",
      "    %1548 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1549 = mhlo.multiply %1548, %1547 : tensor<f32>\n",
      "    %1550 = mhlo.cosine %1549 : tensor<f32>\n",
      "    %1551 = mhlo.sine %1549 : tensor<f32>\n",
      "    %1552 = mhlo.negate %1551 : tensor<f32>\n",
      "    %1553 = mhlo.convert %1550 : tensor<f32>\n",
      "    %1554 = mhlo.convert %1552 : tensor<f32>\n",
      "    %1555 = \"mhlo.broadcast_in_dim\"(%1553) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1556 = \"mhlo.broadcast_in_dim\"(%1554) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1557 = \"mhlo.concatenate\"(%1555, %1556) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1558 = mhlo.convert %1551 : tensor<f32>\n",
      "    %1559 = mhlo.convert %1550 : tensor<f32>\n",
      "    %1560 = \"mhlo.broadcast_in_dim\"(%1558) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1561 = \"mhlo.broadcast_in_dim\"(%1559) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1562 = \"mhlo.concatenate\"(%1560, %1561) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1563 = \"mhlo.broadcast_in_dim\"(%1557) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1564 = \"mhlo.broadcast_in_dim\"(%1562) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1565 = \"mhlo.concatenate\"(%1563, %1564) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1566 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1567 = mhlo.constant dense<0.583333313> : tensor<f32>\n",
      "    %1568 = mhlo.convert %1566 : tensor<f32>\n",
      "    %1569 = mhlo.convert %1567 : tensor<f32>\n",
      "    %1570 = mhlo.power %1568, %1569 : tensor<f32>\n",
      "    %1571 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1572 = mhlo.divide %1571, %1570 : tensor<f32>\n",
      "    %1573 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1574 = mhlo.multiply %1573, %1572 : tensor<f32>\n",
      "    %1575 = mhlo.cosine %1574 : tensor<f32>\n",
      "    %1576 = mhlo.sine %1574 : tensor<f32>\n",
      "    %1577 = mhlo.negate %1576 : tensor<f32>\n",
      "    %1578 = mhlo.convert %1575 : tensor<f32>\n",
      "    %1579 = mhlo.convert %1577 : tensor<f32>\n",
      "    %1580 = \"mhlo.broadcast_in_dim\"(%1578) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1581 = \"mhlo.broadcast_in_dim\"(%1579) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1582 = \"mhlo.concatenate\"(%1580, %1581) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1583 = mhlo.convert %1576 : tensor<f32>\n",
      "    %1584 = mhlo.convert %1575 : tensor<f32>\n",
      "    %1585 = \"mhlo.broadcast_in_dim\"(%1583) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1586 = \"mhlo.broadcast_in_dim\"(%1584) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1587 = \"mhlo.concatenate\"(%1585, %1586) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1588 = \"mhlo.broadcast_in_dim\"(%1582) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1589 = \"mhlo.broadcast_in_dim\"(%1587) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1590 = \"mhlo.concatenate\"(%1588, %1589) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1591 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1592 = mhlo.constant dense<6.250000e-01> : tensor<f32>\n",
      "    %1593 = mhlo.convert %1591 : tensor<f32>\n",
      "    %1594 = mhlo.convert %1592 : tensor<f32>\n",
      "    %1595 = mhlo.power %1593, %1594 : tensor<f32>\n",
      "    %1596 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1597 = mhlo.divide %1596, %1595 : tensor<f32>\n",
      "    %1598 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1599 = mhlo.multiply %1598, %1597 : tensor<f32>\n",
      "    %1600 = mhlo.cosine %1599 : tensor<f32>\n",
      "    %1601 = mhlo.sine %1599 : tensor<f32>\n",
      "    %1602 = mhlo.negate %1601 : tensor<f32>\n",
      "    %1603 = mhlo.convert %1600 : tensor<f32>\n",
      "    %1604 = mhlo.convert %1602 : tensor<f32>\n",
      "    %1605 = \"mhlo.broadcast_in_dim\"(%1603) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1606 = \"mhlo.broadcast_in_dim\"(%1604) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1607 = \"mhlo.concatenate\"(%1605, %1606) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1608 = mhlo.convert %1601 : tensor<f32>\n",
      "    %1609 = mhlo.convert %1600 : tensor<f32>\n",
      "    %1610 = \"mhlo.broadcast_in_dim\"(%1608) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1611 = \"mhlo.broadcast_in_dim\"(%1609) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1612 = \"mhlo.concatenate\"(%1610, %1611) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1613 = \"mhlo.broadcast_in_dim\"(%1607) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1614 = \"mhlo.broadcast_in_dim\"(%1612) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1615 = \"mhlo.concatenate\"(%1613, %1614) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1616 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1617 = mhlo.constant dense<0.666666686> : tensor<f32>\n",
      "    %1618 = mhlo.convert %1616 : tensor<f32>\n",
      "    %1619 = mhlo.convert %1617 : tensor<f32>\n",
      "    %1620 = mhlo.power %1618, %1619 : tensor<f32>\n",
      "    %1621 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1622 = mhlo.divide %1621, %1620 : tensor<f32>\n",
      "    %1623 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1624 = mhlo.multiply %1623, %1622 : tensor<f32>\n",
      "    %1625 = mhlo.cosine %1624 : tensor<f32>\n",
      "    %1626 = mhlo.sine %1624 : tensor<f32>\n",
      "    %1627 = mhlo.negate %1626 : tensor<f32>\n",
      "    %1628 = mhlo.convert %1625 : tensor<f32>\n",
      "    %1629 = mhlo.convert %1627 : tensor<f32>\n",
      "    %1630 = \"mhlo.broadcast_in_dim\"(%1628) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1631 = \"mhlo.broadcast_in_dim\"(%1629) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1632 = \"mhlo.concatenate\"(%1630, %1631) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1633 = mhlo.convert %1626 : tensor<f32>\n",
      "    %1634 = mhlo.convert %1625 : tensor<f32>\n",
      "    %1635 = \"mhlo.broadcast_in_dim\"(%1633) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1636 = \"mhlo.broadcast_in_dim\"(%1634) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1637 = \"mhlo.concatenate\"(%1635, %1636) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1638 = \"mhlo.broadcast_in_dim\"(%1632) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1639 = \"mhlo.broadcast_in_dim\"(%1637) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1640 = \"mhlo.concatenate\"(%1638, %1639) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1641 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1642 = mhlo.constant dense<0.708333313> : tensor<f32>\n",
      "    %1643 = mhlo.convert %1641 : tensor<f32>\n",
      "    %1644 = mhlo.convert %1642 : tensor<f32>\n",
      "    %1645 = mhlo.power %1643, %1644 : tensor<f32>\n",
      "    %1646 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1647 = mhlo.divide %1646, %1645 : tensor<f32>\n",
      "    %1648 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1649 = mhlo.multiply %1648, %1647 : tensor<f32>\n",
      "    %1650 = mhlo.cosine %1649 : tensor<f32>\n",
      "    %1651 = mhlo.sine %1649 : tensor<f32>\n",
      "    %1652 = mhlo.negate %1651 : tensor<f32>\n",
      "    %1653 = mhlo.convert %1650 : tensor<f32>\n",
      "    %1654 = mhlo.convert %1652 : tensor<f32>\n",
      "    %1655 = \"mhlo.broadcast_in_dim\"(%1653) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1656 = \"mhlo.broadcast_in_dim\"(%1654) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1657 = \"mhlo.concatenate\"(%1655, %1656) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1658 = mhlo.convert %1651 : tensor<f32>\n",
      "    %1659 = mhlo.convert %1650 : tensor<f32>\n",
      "    %1660 = \"mhlo.broadcast_in_dim\"(%1658) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1661 = \"mhlo.broadcast_in_dim\"(%1659) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1662 = \"mhlo.concatenate\"(%1660, %1661) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1663 = \"mhlo.broadcast_in_dim\"(%1657) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1664 = \"mhlo.broadcast_in_dim\"(%1662) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1665 = \"mhlo.concatenate\"(%1663, %1664) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1666 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1667 = mhlo.constant dense<7.500000e-01> : tensor<f32>\n",
      "    %1668 = mhlo.convert %1666 : tensor<f32>\n",
      "    %1669 = mhlo.convert %1667 : tensor<f32>\n",
      "    %1670 = mhlo.power %1668, %1669 : tensor<f32>\n",
      "    %1671 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1672 = mhlo.divide %1671, %1670 : tensor<f32>\n",
      "    %1673 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1674 = mhlo.multiply %1673, %1672 : tensor<f32>\n",
      "    %1675 = mhlo.cosine %1674 : tensor<f32>\n",
      "    %1676 = mhlo.sine %1674 : tensor<f32>\n",
      "    %1677 = mhlo.negate %1676 : tensor<f32>\n",
      "    %1678 = mhlo.convert %1675 : tensor<f32>\n",
      "    %1679 = mhlo.convert %1677 : tensor<f32>\n",
      "    %1680 = \"mhlo.broadcast_in_dim\"(%1678) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1681 = \"mhlo.broadcast_in_dim\"(%1679) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1682 = \"mhlo.concatenate\"(%1680, %1681) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1683 = mhlo.convert %1676 : tensor<f32>\n",
      "    %1684 = mhlo.convert %1675 : tensor<f32>\n",
      "    %1685 = \"mhlo.broadcast_in_dim\"(%1683) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1686 = \"mhlo.broadcast_in_dim\"(%1684) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1687 = \"mhlo.concatenate\"(%1685, %1686) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1688 = \"mhlo.broadcast_in_dim\"(%1682) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1689 = \"mhlo.broadcast_in_dim\"(%1687) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1690 = \"mhlo.concatenate\"(%1688, %1689) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1691 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1692 = mhlo.constant dense<0.791666686> : tensor<f32>\n",
      "    %1693 = mhlo.convert %1691 : tensor<f32>\n",
      "    %1694 = mhlo.convert %1692 : tensor<f32>\n",
      "    %1695 = mhlo.power %1693, %1694 : tensor<f32>\n",
      "    %1696 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1697 = mhlo.divide %1696, %1695 : tensor<f32>\n",
      "    %1698 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1699 = mhlo.multiply %1698, %1697 : tensor<f32>\n",
      "    %1700 = mhlo.cosine %1699 : tensor<f32>\n",
      "    %1701 = mhlo.sine %1699 : tensor<f32>\n",
      "    %1702 = mhlo.negate %1701 : tensor<f32>\n",
      "    %1703 = mhlo.convert %1700 : tensor<f32>\n",
      "    %1704 = mhlo.convert %1702 : tensor<f32>\n",
      "    %1705 = \"mhlo.broadcast_in_dim\"(%1703) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1706 = \"mhlo.broadcast_in_dim\"(%1704) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1707 = \"mhlo.concatenate\"(%1705, %1706) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1708 = mhlo.convert %1701 : tensor<f32>\n",
      "    %1709 = mhlo.convert %1700 : tensor<f32>\n",
      "    %1710 = \"mhlo.broadcast_in_dim\"(%1708) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1711 = \"mhlo.broadcast_in_dim\"(%1709) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1712 = \"mhlo.concatenate\"(%1710, %1711) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1713 = \"mhlo.broadcast_in_dim\"(%1707) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1714 = \"mhlo.broadcast_in_dim\"(%1712) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1715 = \"mhlo.concatenate\"(%1713, %1714) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1716 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1717 = mhlo.constant dense<0.833333313> : tensor<f32>\n",
      "    %1718 = mhlo.convert %1716 : tensor<f32>\n",
      "    %1719 = mhlo.convert %1717 : tensor<f32>\n",
      "    %1720 = mhlo.power %1718, %1719 : tensor<f32>\n",
      "    %1721 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1722 = mhlo.divide %1721, %1720 : tensor<f32>\n",
      "    %1723 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1724 = mhlo.multiply %1723, %1722 : tensor<f32>\n",
      "    %1725 = mhlo.cosine %1724 : tensor<f32>\n",
      "    %1726 = mhlo.sine %1724 : tensor<f32>\n",
      "    %1727 = mhlo.negate %1726 : tensor<f32>\n",
      "    %1728 = mhlo.convert %1725 : tensor<f32>\n",
      "    %1729 = mhlo.convert %1727 : tensor<f32>\n",
      "    %1730 = \"mhlo.broadcast_in_dim\"(%1728) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1731 = \"mhlo.broadcast_in_dim\"(%1729) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1732 = \"mhlo.concatenate\"(%1730, %1731) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1733 = mhlo.convert %1726 : tensor<f32>\n",
      "    %1734 = mhlo.convert %1725 : tensor<f32>\n",
      "    %1735 = \"mhlo.broadcast_in_dim\"(%1733) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1736 = \"mhlo.broadcast_in_dim\"(%1734) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1737 = \"mhlo.concatenate\"(%1735, %1736) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1738 = \"mhlo.broadcast_in_dim\"(%1732) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1739 = \"mhlo.broadcast_in_dim\"(%1737) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1740 = \"mhlo.concatenate\"(%1738, %1739) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1741 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1742 = mhlo.constant dense<8.750000e-01> : tensor<f32>\n",
      "    %1743 = mhlo.convert %1741 : tensor<f32>\n",
      "    %1744 = mhlo.convert %1742 : tensor<f32>\n",
      "    %1745 = mhlo.power %1743, %1744 : tensor<f32>\n",
      "    %1746 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1747 = mhlo.divide %1746, %1745 : tensor<f32>\n",
      "    %1748 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1749 = mhlo.multiply %1748, %1747 : tensor<f32>\n",
      "    %1750 = mhlo.cosine %1749 : tensor<f32>\n",
      "    %1751 = mhlo.sine %1749 : tensor<f32>\n",
      "    %1752 = mhlo.negate %1751 : tensor<f32>\n",
      "    %1753 = mhlo.convert %1750 : tensor<f32>\n",
      "    %1754 = mhlo.convert %1752 : tensor<f32>\n",
      "    %1755 = \"mhlo.broadcast_in_dim\"(%1753) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1756 = \"mhlo.broadcast_in_dim\"(%1754) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1757 = \"mhlo.concatenate\"(%1755, %1756) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1758 = mhlo.convert %1751 : tensor<f32>\n",
      "    %1759 = mhlo.convert %1750 : tensor<f32>\n",
      "    %1760 = \"mhlo.broadcast_in_dim\"(%1758) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1761 = \"mhlo.broadcast_in_dim\"(%1759) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1762 = \"mhlo.concatenate\"(%1760, %1761) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1763 = \"mhlo.broadcast_in_dim\"(%1757) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1764 = \"mhlo.broadcast_in_dim\"(%1762) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1765 = \"mhlo.concatenate\"(%1763, %1764) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1766 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1767 = mhlo.constant dense<0.916666686> : tensor<f32>\n",
      "    %1768 = mhlo.convert %1766 : tensor<f32>\n",
      "    %1769 = mhlo.convert %1767 : tensor<f32>\n",
      "    %1770 = mhlo.power %1768, %1769 : tensor<f32>\n",
      "    %1771 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1772 = mhlo.divide %1771, %1770 : tensor<f32>\n",
      "    %1773 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1774 = mhlo.multiply %1773, %1772 : tensor<f32>\n",
      "    %1775 = mhlo.cosine %1774 : tensor<f32>\n",
      "    %1776 = mhlo.sine %1774 : tensor<f32>\n",
      "    %1777 = mhlo.negate %1776 : tensor<f32>\n",
      "    %1778 = mhlo.convert %1775 : tensor<f32>\n",
      "    %1779 = mhlo.convert %1777 : tensor<f32>\n",
      "    %1780 = \"mhlo.broadcast_in_dim\"(%1778) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1781 = \"mhlo.broadcast_in_dim\"(%1779) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1782 = \"mhlo.concatenate\"(%1780, %1781) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1783 = mhlo.convert %1776 : tensor<f32>\n",
      "    %1784 = mhlo.convert %1775 : tensor<f32>\n",
      "    %1785 = \"mhlo.broadcast_in_dim\"(%1783) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1786 = \"mhlo.broadcast_in_dim\"(%1784) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1787 = \"mhlo.concatenate\"(%1785, %1786) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1788 = \"mhlo.broadcast_in_dim\"(%1782) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1789 = \"mhlo.broadcast_in_dim\"(%1787) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1790 = \"mhlo.concatenate\"(%1788, %1789) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1791 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1792 = mhlo.constant dense<0.958333313> : tensor<f32>\n",
      "    %1793 = mhlo.convert %1791 : tensor<f32>\n",
      "    %1794 = mhlo.convert %1792 : tensor<f32>\n",
      "    %1795 = mhlo.power %1793, %1794 : tensor<f32>\n",
      "    %1796 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1797 = mhlo.divide %1796, %1795 : tensor<f32>\n",
      "    %1798 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1799 = mhlo.multiply %1798, %1797 : tensor<f32>\n",
      "    %1800 = mhlo.cosine %1799 : tensor<f32>\n",
      "    %1801 = mhlo.sine %1799 : tensor<f32>\n",
      "    %1802 = mhlo.negate %1801 : tensor<f32>\n",
      "    %1803 = mhlo.convert %1800 : tensor<f32>\n",
      "    %1804 = mhlo.convert %1802 : tensor<f32>\n",
      "    %1805 = \"mhlo.broadcast_in_dim\"(%1803) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1806 = \"mhlo.broadcast_in_dim\"(%1804) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1807 = \"mhlo.concatenate\"(%1805, %1806) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1808 = mhlo.convert %1801 : tensor<f32>\n",
      "    %1809 = mhlo.convert %1800 : tensor<f32>\n",
      "    %1810 = \"mhlo.broadcast_in_dim\"(%1808) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1811 = \"mhlo.broadcast_in_dim\"(%1809) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1812 = \"mhlo.concatenate\"(%1810, %1811) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1813 = \"mhlo.broadcast_in_dim\"(%1807) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1814 = \"mhlo.broadcast_in_dim\"(%1812) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1815 = \"mhlo.concatenate\"(%1813, %1814) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1816 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1817 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1818 = mhlo.convert %1816 : tensor<f32>\n",
      "    %1819 = mhlo.convert %1817 : tensor<f32>\n",
      "    %1820 = mhlo.power %1818, %1819 : tensor<f32>\n",
      "    %1821 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1822 = mhlo.divide %1821, %1820 : tensor<f32>\n",
      "    %1823 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1824 = mhlo.multiply %1823, %1822 : tensor<f32>\n",
      "    %1825 = mhlo.cosine %1824 : tensor<f32>\n",
      "    %1826 = mhlo.sine %1824 : tensor<f32>\n",
      "    %1827 = mhlo.negate %1826 : tensor<f32>\n",
      "    %1828 = mhlo.convert %1825 : tensor<f32>\n",
      "    %1829 = mhlo.convert %1827 : tensor<f32>\n",
      "    %1830 = \"mhlo.broadcast_in_dim\"(%1828) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1831 = \"mhlo.broadcast_in_dim\"(%1829) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1832 = \"mhlo.concatenate\"(%1830, %1831) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1833 = mhlo.convert %1826 : tensor<f32>\n",
      "    %1834 = mhlo.convert %1825 : tensor<f32>\n",
      "    %1835 = \"mhlo.broadcast_in_dim\"(%1833) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1836 = \"mhlo.broadcast_in_dim\"(%1834) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1837 = \"mhlo.concatenate\"(%1835, %1836) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1838 = \"mhlo.broadcast_in_dim\"(%1832) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1839 = \"mhlo.broadcast_in_dim\"(%1837) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1840 = \"mhlo.concatenate\"(%1838, %1839) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1841 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1842 = mhlo.constant dense<0.0416666679> : tensor<f32>\n",
      "    %1843 = mhlo.convert %1841 : tensor<f32>\n",
      "    %1844 = mhlo.convert %1842 : tensor<f32>\n",
      "    %1845 = mhlo.power %1843, %1844 : tensor<f32>\n",
      "    %1846 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1847 = mhlo.divide %1846, %1845 : tensor<f32>\n",
      "    %1848 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1849 = mhlo.multiply %1848, %1847 : tensor<f32>\n",
      "    %1850 = mhlo.cosine %1849 : tensor<f32>\n",
      "    %1851 = mhlo.sine %1849 : tensor<f32>\n",
      "    %1852 = mhlo.negate %1851 : tensor<f32>\n",
      "    %1853 = mhlo.convert %1850 : tensor<f32>\n",
      "    %1854 = mhlo.convert %1852 : tensor<f32>\n",
      "    %1855 = \"mhlo.broadcast_in_dim\"(%1853) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1856 = \"mhlo.broadcast_in_dim\"(%1854) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1857 = \"mhlo.concatenate\"(%1855, %1856) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1858 = mhlo.convert %1851 : tensor<f32>\n",
      "    %1859 = mhlo.convert %1850 : tensor<f32>\n",
      "    %1860 = \"mhlo.broadcast_in_dim\"(%1858) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1861 = \"mhlo.broadcast_in_dim\"(%1859) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1862 = \"mhlo.concatenate\"(%1860, %1861) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1863 = \"mhlo.broadcast_in_dim\"(%1857) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1864 = \"mhlo.broadcast_in_dim\"(%1862) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1865 = \"mhlo.concatenate\"(%1863, %1864) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1866 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1867 = mhlo.constant dense<0.0833333358> : tensor<f32>\n",
      "    %1868 = mhlo.convert %1866 : tensor<f32>\n",
      "    %1869 = mhlo.convert %1867 : tensor<f32>\n",
      "    %1870 = mhlo.power %1868, %1869 : tensor<f32>\n",
      "    %1871 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1872 = mhlo.divide %1871, %1870 : tensor<f32>\n",
      "    %1873 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1874 = mhlo.multiply %1873, %1872 : tensor<f32>\n",
      "    %1875 = mhlo.cosine %1874 : tensor<f32>\n",
      "    %1876 = mhlo.sine %1874 : tensor<f32>\n",
      "    %1877 = mhlo.negate %1876 : tensor<f32>\n",
      "    %1878 = mhlo.convert %1875 : tensor<f32>\n",
      "    %1879 = mhlo.convert %1877 : tensor<f32>\n",
      "    %1880 = \"mhlo.broadcast_in_dim\"(%1878) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1881 = \"mhlo.broadcast_in_dim\"(%1879) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1882 = \"mhlo.concatenate\"(%1880, %1881) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1883 = mhlo.convert %1876 : tensor<f32>\n",
      "    %1884 = mhlo.convert %1875 : tensor<f32>\n",
      "    %1885 = \"mhlo.broadcast_in_dim\"(%1883) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1886 = \"mhlo.broadcast_in_dim\"(%1884) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1887 = \"mhlo.concatenate\"(%1885, %1886) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1888 = \"mhlo.broadcast_in_dim\"(%1882) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1889 = \"mhlo.broadcast_in_dim\"(%1887) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1890 = \"mhlo.concatenate\"(%1888, %1889) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1891 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1892 = mhlo.constant dense<1.250000e-01> : tensor<f32>\n",
      "    %1893 = mhlo.convert %1891 : tensor<f32>\n",
      "    %1894 = mhlo.convert %1892 : tensor<f32>\n",
      "    %1895 = mhlo.power %1893, %1894 : tensor<f32>\n",
      "    %1896 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1897 = mhlo.divide %1896, %1895 : tensor<f32>\n",
      "    %1898 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1899 = mhlo.multiply %1898, %1897 : tensor<f32>\n",
      "    %1900 = mhlo.cosine %1899 : tensor<f32>\n",
      "    %1901 = mhlo.sine %1899 : tensor<f32>\n",
      "    %1902 = mhlo.negate %1901 : tensor<f32>\n",
      "    %1903 = mhlo.convert %1900 : tensor<f32>\n",
      "    %1904 = mhlo.convert %1902 : tensor<f32>\n",
      "    %1905 = \"mhlo.broadcast_in_dim\"(%1903) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1906 = \"mhlo.broadcast_in_dim\"(%1904) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1907 = \"mhlo.concatenate\"(%1905, %1906) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1908 = mhlo.convert %1901 : tensor<f32>\n",
      "    %1909 = mhlo.convert %1900 : tensor<f32>\n",
      "    %1910 = \"mhlo.broadcast_in_dim\"(%1908) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1911 = \"mhlo.broadcast_in_dim\"(%1909) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1912 = \"mhlo.concatenate\"(%1910, %1911) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1913 = \"mhlo.broadcast_in_dim\"(%1907) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1914 = \"mhlo.broadcast_in_dim\"(%1912) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1915 = \"mhlo.concatenate\"(%1913, %1914) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1916 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1917 = mhlo.constant dense<0.166666672> : tensor<f32>\n",
      "    %1918 = mhlo.convert %1916 : tensor<f32>\n",
      "    %1919 = mhlo.convert %1917 : tensor<f32>\n",
      "    %1920 = mhlo.power %1918, %1919 : tensor<f32>\n",
      "    %1921 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1922 = mhlo.divide %1921, %1920 : tensor<f32>\n",
      "    %1923 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1924 = mhlo.multiply %1923, %1922 : tensor<f32>\n",
      "    %1925 = mhlo.cosine %1924 : tensor<f32>\n",
      "    %1926 = mhlo.sine %1924 : tensor<f32>\n",
      "    %1927 = mhlo.negate %1926 : tensor<f32>\n",
      "    %1928 = mhlo.convert %1925 : tensor<f32>\n",
      "    %1929 = mhlo.convert %1927 : tensor<f32>\n",
      "    %1930 = \"mhlo.broadcast_in_dim\"(%1928) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1931 = \"mhlo.broadcast_in_dim\"(%1929) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1932 = \"mhlo.concatenate\"(%1930, %1931) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1933 = mhlo.convert %1926 : tensor<f32>\n",
      "    %1934 = mhlo.convert %1925 : tensor<f32>\n",
      "    %1935 = \"mhlo.broadcast_in_dim\"(%1933) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1936 = \"mhlo.broadcast_in_dim\"(%1934) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1937 = \"mhlo.concatenate\"(%1935, %1936) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1938 = \"mhlo.broadcast_in_dim\"(%1932) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1939 = \"mhlo.broadcast_in_dim\"(%1937) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1940 = \"mhlo.concatenate\"(%1938, %1939) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1941 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1942 = mhlo.constant dense<0.208333328> : tensor<f32>\n",
      "    %1943 = mhlo.convert %1941 : tensor<f32>\n",
      "    %1944 = mhlo.convert %1942 : tensor<f32>\n",
      "    %1945 = mhlo.power %1943, %1944 : tensor<f32>\n",
      "    %1946 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1947 = mhlo.divide %1946, %1945 : tensor<f32>\n",
      "    %1948 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1949 = mhlo.multiply %1948, %1947 : tensor<f32>\n",
      "    %1950 = mhlo.cosine %1949 : tensor<f32>\n",
      "    %1951 = mhlo.sine %1949 : tensor<f32>\n",
      "    %1952 = mhlo.negate %1951 : tensor<f32>\n",
      "    %1953 = mhlo.convert %1950 : tensor<f32>\n",
      "    %1954 = mhlo.convert %1952 : tensor<f32>\n",
      "    %1955 = \"mhlo.broadcast_in_dim\"(%1953) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1956 = \"mhlo.broadcast_in_dim\"(%1954) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1957 = \"mhlo.concatenate\"(%1955, %1956) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1958 = mhlo.convert %1951 : tensor<f32>\n",
      "    %1959 = mhlo.convert %1950 : tensor<f32>\n",
      "    %1960 = \"mhlo.broadcast_in_dim\"(%1958) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1961 = \"mhlo.broadcast_in_dim\"(%1959) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1962 = \"mhlo.concatenate\"(%1960, %1961) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1963 = \"mhlo.broadcast_in_dim\"(%1957) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1964 = \"mhlo.broadcast_in_dim\"(%1962) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1965 = \"mhlo.concatenate\"(%1963, %1964) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1966 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1967 = mhlo.constant dense<2.500000e-01> : tensor<f32>\n",
      "    %1968 = mhlo.convert %1966 : tensor<f32>\n",
      "    %1969 = mhlo.convert %1967 : tensor<f32>\n",
      "    %1970 = mhlo.power %1968, %1969 : tensor<f32>\n",
      "    %1971 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1972 = mhlo.divide %1971, %1970 : tensor<f32>\n",
      "    %1973 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1974 = mhlo.multiply %1973, %1972 : tensor<f32>\n",
      "    %1975 = mhlo.cosine %1974 : tensor<f32>\n",
      "    %1976 = mhlo.sine %1974 : tensor<f32>\n",
      "    %1977 = mhlo.negate %1976 : tensor<f32>\n",
      "    %1978 = mhlo.convert %1975 : tensor<f32>\n",
      "    %1979 = mhlo.convert %1977 : tensor<f32>\n",
      "    %1980 = \"mhlo.broadcast_in_dim\"(%1978) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1981 = \"mhlo.broadcast_in_dim\"(%1979) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1982 = \"mhlo.concatenate\"(%1980, %1981) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1983 = mhlo.convert %1976 : tensor<f32>\n",
      "    %1984 = mhlo.convert %1975 : tensor<f32>\n",
      "    %1985 = \"mhlo.broadcast_in_dim\"(%1983) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1986 = \"mhlo.broadcast_in_dim\"(%1984) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %1987 = \"mhlo.concatenate\"(%1985, %1986) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %1988 = \"mhlo.broadcast_in_dim\"(%1982) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1989 = \"mhlo.broadcast_in_dim\"(%1987) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %1990 = \"mhlo.concatenate\"(%1988, %1989) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %1991 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %1992 = mhlo.constant dense<0.291666657> : tensor<f32>\n",
      "    %1993 = mhlo.convert %1991 : tensor<f32>\n",
      "    %1994 = mhlo.convert %1992 : tensor<f32>\n",
      "    %1995 = mhlo.power %1993, %1994 : tensor<f32>\n",
      "    %1996 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %1997 = mhlo.divide %1996, %1995 : tensor<f32>\n",
      "    %1998 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %1999 = mhlo.multiply %1998, %1997 : tensor<f32>\n",
      "    %2000 = mhlo.cosine %1999 : tensor<f32>\n",
      "    %2001 = mhlo.sine %1999 : tensor<f32>\n",
      "    %2002 = mhlo.negate %2001 : tensor<f32>\n",
      "    %2003 = mhlo.convert %2000 : tensor<f32>\n",
      "    %2004 = mhlo.convert %2002 : tensor<f32>\n",
      "    %2005 = \"mhlo.broadcast_in_dim\"(%2003) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2006 = \"mhlo.broadcast_in_dim\"(%2004) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2007 = \"mhlo.concatenate\"(%2005, %2006) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2008 = mhlo.convert %2001 : tensor<f32>\n",
      "    %2009 = mhlo.convert %2000 : tensor<f32>\n",
      "    %2010 = \"mhlo.broadcast_in_dim\"(%2008) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2011 = \"mhlo.broadcast_in_dim\"(%2009) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2012 = \"mhlo.concatenate\"(%2010, %2011) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2013 = \"mhlo.broadcast_in_dim\"(%2007) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2014 = \"mhlo.broadcast_in_dim\"(%2012) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2015 = \"mhlo.concatenate\"(%2013, %2014) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2016 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2017 = mhlo.constant dense<0.333333343> : tensor<f32>\n",
      "    %2018 = mhlo.convert %2016 : tensor<f32>\n",
      "    %2019 = mhlo.convert %2017 : tensor<f32>\n",
      "    %2020 = mhlo.power %2018, %2019 : tensor<f32>\n",
      "    %2021 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2022 = mhlo.divide %2021, %2020 : tensor<f32>\n",
      "    %2023 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2024 = mhlo.multiply %2023, %2022 : tensor<f32>\n",
      "    %2025 = mhlo.cosine %2024 : tensor<f32>\n",
      "    %2026 = mhlo.sine %2024 : tensor<f32>\n",
      "    %2027 = mhlo.negate %2026 : tensor<f32>\n",
      "    %2028 = mhlo.convert %2025 : tensor<f32>\n",
      "    %2029 = mhlo.convert %2027 : tensor<f32>\n",
      "    %2030 = \"mhlo.broadcast_in_dim\"(%2028) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2031 = \"mhlo.broadcast_in_dim\"(%2029) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2032 = \"mhlo.concatenate\"(%2030, %2031) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2033 = mhlo.convert %2026 : tensor<f32>\n",
      "    %2034 = mhlo.convert %2025 : tensor<f32>\n",
      "    %2035 = \"mhlo.broadcast_in_dim\"(%2033) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2036 = \"mhlo.broadcast_in_dim\"(%2034) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2037 = \"mhlo.concatenate\"(%2035, %2036) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2038 = \"mhlo.broadcast_in_dim\"(%2032) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2039 = \"mhlo.broadcast_in_dim\"(%2037) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2040 = \"mhlo.concatenate\"(%2038, %2039) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2041 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2042 = mhlo.constant dense<3.750000e-01> : tensor<f32>\n",
      "    %2043 = mhlo.convert %2041 : tensor<f32>\n",
      "    %2044 = mhlo.convert %2042 : tensor<f32>\n",
      "    %2045 = mhlo.power %2043, %2044 : tensor<f32>\n",
      "    %2046 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2047 = mhlo.divide %2046, %2045 : tensor<f32>\n",
      "    %2048 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2049 = mhlo.multiply %2048, %2047 : tensor<f32>\n",
      "    %2050 = mhlo.cosine %2049 : tensor<f32>\n",
      "    %2051 = mhlo.sine %2049 : tensor<f32>\n",
      "    %2052 = mhlo.negate %2051 : tensor<f32>\n",
      "    %2053 = mhlo.convert %2050 : tensor<f32>\n",
      "    %2054 = mhlo.convert %2052 : tensor<f32>\n",
      "    %2055 = \"mhlo.broadcast_in_dim\"(%2053) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2056 = \"mhlo.broadcast_in_dim\"(%2054) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2057 = \"mhlo.concatenate\"(%2055, %2056) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2058 = mhlo.convert %2051 : tensor<f32>\n",
      "    %2059 = mhlo.convert %2050 : tensor<f32>\n",
      "    %2060 = \"mhlo.broadcast_in_dim\"(%2058) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2061 = \"mhlo.broadcast_in_dim\"(%2059) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2062 = \"mhlo.concatenate\"(%2060, %2061) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2063 = \"mhlo.broadcast_in_dim\"(%2057) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2064 = \"mhlo.broadcast_in_dim\"(%2062) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2065 = \"mhlo.concatenate\"(%2063, %2064) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2066 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2067 = mhlo.constant dense<0.416666657> : tensor<f32>\n",
      "    %2068 = mhlo.convert %2066 : tensor<f32>\n",
      "    %2069 = mhlo.convert %2067 : tensor<f32>\n",
      "    %2070 = mhlo.power %2068, %2069 : tensor<f32>\n",
      "    %2071 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2072 = mhlo.divide %2071, %2070 : tensor<f32>\n",
      "    %2073 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2074 = mhlo.multiply %2073, %2072 : tensor<f32>\n",
      "    %2075 = mhlo.cosine %2074 : tensor<f32>\n",
      "    %2076 = mhlo.sine %2074 : tensor<f32>\n",
      "    %2077 = mhlo.negate %2076 : tensor<f32>\n",
      "    %2078 = mhlo.convert %2075 : tensor<f32>\n",
      "    %2079 = mhlo.convert %2077 : tensor<f32>\n",
      "    %2080 = \"mhlo.broadcast_in_dim\"(%2078) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2081 = \"mhlo.broadcast_in_dim\"(%2079) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2082 = \"mhlo.concatenate\"(%2080, %2081) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2083 = mhlo.convert %2076 : tensor<f32>\n",
      "    %2084 = mhlo.convert %2075 : tensor<f32>\n",
      "    %2085 = \"mhlo.broadcast_in_dim\"(%2083) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2086 = \"mhlo.broadcast_in_dim\"(%2084) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2087 = \"mhlo.concatenate\"(%2085, %2086) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2088 = \"mhlo.broadcast_in_dim\"(%2082) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2089 = \"mhlo.broadcast_in_dim\"(%2087) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2090 = \"mhlo.concatenate\"(%2088, %2089) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2091 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2092 = mhlo.constant dense<0.458333343> : tensor<f32>\n",
      "    %2093 = mhlo.convert %2091 : tensor<f32>\n",
      "    %2094 = mhlo.convert %2092 : tensor<f32>\n",
      "    %2095 = mhlo.power %2093, %2094 : tensor<f32>\n",
      "    %2096 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2097 = mhlo.divide %2096, %2095 : tensor<f32>\n",
      "    %2098 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2099 = mhlo.multiply %2098, %2097 : tensor<f32>\n",
      "    %2100 = mhlo.cosine %2099 : tensor<f32>\n",
      "    %2101 = mhlo.sine %2099 : tensor<f32>\n",
      "    %2102 = mhlo.negate %2101 : tensor<f32>\n",
      "    %2103 = mhlo.convert %2100 : tensor<f32>\n",
      "    %2104 = mhlo.convert %2102 : tensor<f32>\n",
      "    %2105 = \"mhlo.broadcast_in_dim\"(%2103) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2106 = \"mhlo.broadcast_in_dim\"(%2104) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2107 = \"mhlo.concatenate\"(%2105, %2106) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2108 = mhlo.convert %2101 : tensor<f32>\n",
      "    %2109 = mhlo.convert %2100 : tensor<f32>\n",
      "    %2110 = \"mhlo.broadcast_in_dim\"(%2108) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2111 = \"mhlo.broadcast_in_dim\"(%2109) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2112 = \"mhlo.concatenate\"(%2110, %2111) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2113 = \"mhlo.broadcast_in_dim\"(%2107) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2114 = \"mhlo.broadcast_in_dim\"(%2112) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2115 = \"mhlo.concatenate\"(%2113, %2114) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2116 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2117 = mhlo.constant dense<5.000000e-01> : tensor<f32>\n",
      "    %2118 = mhlo.convert %2116 : tensor<f32>\n",
      "    %2119 = mhlo.convert %2117 : tensor<f32>\n",
      "    %2120 = mhlo.power %2118, %2119 : tensor<f32>\n",
      "    %2121 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2122 = mhlo.divide %2121, %2120 : tensor<f32>\n",
      "    %2123 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2124 = mhlo.multiply %2123, %2122 : tensor<f32>\n",
      "    %2125 = mhlo.cosine %2124 : tensor<f32>\n",
      "    %2126 = mhlo.sine %2124 : tensor<f32>\n",
      "    %2127 = mhlo.negate %2126 : tensor<f32>\n",
      "    %2128 = mhlo.convert %2125 : tensor<f32>\n",
      "    %2129 = mhlo.convert %2127 : tensor<f32>\n",
      "    %2130 = \"mhlo.broadcast_in_dim\"(%2128) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2131 = \"mhlo.broadcast_in_dim\"(%2129) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2132 = \"mhlo.concatenate\"(%2130, %2131) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2133 = mhlo.convert %2126 : tensor<f32>\n",
      "    %2134 = mhlo.convert %2125 : tensor<f32>\n",
      "    %2135 = \"mhlo.broadcast_in_dim\"(%2133) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2136 = \"mhlo.broadcast_in_dim\"(%2134) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2137 = \"mhlo.concatenate\"(%2135, %2136) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2138 = \"mhlo.broadcast_in_dim\"(%2132) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2139 = \"mhlo.broadcast_in_dim\"(%2137) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2140 = \"mhlo.concatenate\"(%2138, %2139) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2141 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2142 = mhlo.constant dense<0.541666687> : tensor<f32>\n",
      "    %2143 = mhlo.convert %2141 : tensor<f32>\n",
      "    %2144 = mhlo.convert %2142 : tensor<f32>\n",
      "    %2145 = mhlo.power %2143, %2144 : tensor<f32>\n",
      "    %2146 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2147 = mhlo.divide %2146, %2145 : tensor<f32>\n",
      "    %2148 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2149 = mhlo.multiply %2148, %2147 : tensor<f32>\n",
      "    %2150 = mhlo.cosine %2149 : tensor<f32>\n",
      "    %2151 = mhlo.sine %2149 : tensor<f32>\n",
      "    %2152 = mhlo.negate %2151 : tensor<f32>\n",
      "    %2153 = mhlo.convert %2150 : tensor<f32>\n",
      "    %2154 = mhlo.convert %2152 : tensor<f32>\n",
      "    %2155 = \"mhlo.broadcast_in_dim\"(%2153) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2156 = \"mhlo.broadcast_in_dim\"(%2154) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2157 = \"mhlo.concatenate\"(%2155, %2156) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2158 = mhlo.convert %2151 : tensor<f32>\n",
      "    %2159 = mhlo.convert %2150 : tensor<f32>\n",
      "    %2160 = \"mhlo.broadcast_in_dim\"(%2158) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2161 = \"mhlo.broadcast_in_dim\"(%2159) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2162 = \"mhlo.concatenate\"(%2160, %2161) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2163 = \"mhlo.broadcast_in_dim\"(%2157) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2164 = \"mhlo.broadcast_in_dim\"(%2162) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2165 = \"mhlo.concatenate\"(%2163, %2164) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2166 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2167 = mhlo.constant dense<0.583333313> : tensor<f32>\n",
      "    %2168 = mhlo.convert %2166 : tensor<f32>\n",
      "    %2169 = mhlo.convert %2167 : tensor<f32>\n",
      "    %2170 = mhlo.power %2168, %2169 : tensor<f32>\n",
      "    %2171 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2172 = mhlo.divide %2171, %2170 : tensor<f32>\n",
      "    %2173 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2174 = mhlo.multiply %2173, %2172 : tensor<f32>\n",
      "    %2175 = mhlo.cosine %2174 : tensor<f32>\n",
      "    %2176 = mhlo.sine %2174 : tensor<f32>\n",
      "    %2177 = mhlo.negate %2176 : tensor<f32>\n",
      "    %2178 = mhlo.convert %2175 : tensor<f32>\n",
      "    %2179 = mhlo.convert %2177 : tensor<f32>\n",
      "    %2180 = \"mhlo.broadcast_in_dim\"(%2178) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2181 = \"mhlo.broadcast_in_dim\"(%2179) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2182 = \"mhlo.concatenate\"(%2180, %2181) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2183 = mhlo.convert %2176 : tensor<f32>\n",
      "    %2184 = mhlo.convert %2175 : tensor<f32>\n",
      "    %2185 = \"mhlo.broadcast_in_dim\"(%2183) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2186 = \"mhlo.broadcast_in_dim\"(%2184) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2187 = \"mhlo.concatenate\"(%2185, %2186) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2188 = \"mhlo.broadcast_in_dim\"(%2182) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2189 = \"mhlo.broadcast_in_dim\"(%2187) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2190 = \"mhlo.concatenate\"(%2188, %2189) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2191 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2192 = mhlo.constant dense<6.250000e-01> : tensor<f32>\n",
      "    %2193 = mhlo.convert %2191 : tensor<f32>\n",
      "    %2194 = mhlo.convert %2192 : tensor<f32>\n",
      "    %2195 = mhlo.power %2193, %2194 : tensor<f32>\n",
      "    %2196 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2197 = mhlo.divide %2196, %2195 : tensor<f32>\n",
      "    %2198 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2199 = mhlo.multiply %2198, %2197 : tensor<f32>\n",
      "    %2200 = mhlo.cosine %2199 : tensor<f32>\n",
      "    %2201 = mhlo.sine %2199 : tensor<f32>\n",
      "    %2202 = mhlo.negate %2201 : tensor<f32>\n",
      "    %2203 = mhlo.convert %2200 : tensor<f32>\n",
      "    %2204 = mhlo.convert %2202 : tensor<f32>\n",
      "    %2205 = \"mhlo.broadcast_in_dim\"(%2203) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2206 = \"mhlo.broadcast_in_dim\"(%2204) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2207 = \"mhlo.concatenate\"(%2205, %2206) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2208 = mhlo.convert %2201 : tensor<f32>\n",
      "    %2209 = mhlo.convert %2200 : tensor<f32>\n",
      "    %2210 = \"mhlo.broadcast_in_dim\"(%2208) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2211 = \"mhlo.broadcast_in_dim\"(%2209) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2212 = \"mhlo.concatenate\"(%2210, %2211) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2213 = \"mhlo.broadcast_in_dim\"(%2207) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2214 = \"mhlo.broadcast_in_dim\"(%2212) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2215 = \"mhlo.concatenate\"(%2213, %2214) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2216 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2217 = mhlo.constant dense<0.666666686> : tensor<f32>\n",
      "    %2218 = mhlo.convert %2216 : tensor<f32>\n",
      "    %2219 = mhlo.convert %2217 : tensor<f32>\n",
      "    %2220 = mhlo.power %2218, %2219 : tensor<f32>\n",
      "    %2221 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2222 = mhlo.divide %2221, %2220 : tensor<f32>\n",
      "    %2223 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2224 = mhlo.multiply %2223, %2222 : tensor<f32>\n",
      "    %2225 = mhlo.cosine %2224 : tensor<f32>\n",
      "    %2226 = mhlo.sine %2224 : tensor<f32>\n",
      "    %2227 = mhlo.negate %2226 : tensor<f32>\n",
      "    %2228 = mhlo.convert %2225 : tensor<f32>\n",
      "    %2229 = mhlo.convert %2227 : tensor<f32>\n",
      "    %2230 = \"mhlo.broadcast_in_dim\"(%2228) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2231 = \"mhlo.broadcast_in_dim\"(%2229) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2232 = \"mhlo.concatenate\"(%2230, %2231) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2233 = mhlo.convert %2226 : tensor<f32>\n",
      "    %2234 = mhlo.convert %2225 : tensor<f32>\n",
      "    %2235 = \"mhlo.broadcast_in_dim\"(%2233) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2236 = \"mhlo.broadcast_in_dim\"(%2234) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2237 = \"mhlo.concatenate\"(%2235, %2236) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2238 = \"mhlo.broadcast_in_dim\"(%2232) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2239 = \"mhlo.broadcast_in_dim\"(%2237) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2240 = \"mhlo.concatenate\"(%2238, %2239) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2241 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2242 = mhlo.constant dense<0.708333313> : tensor<f32>\n",
      "    %2243 = mhlo.convert %2241 : tensor<f32>\n",
      "    %2244 = mhlo.convert %2242 : tensor<f32>\n",
      "    %2245 = mhlo.power %2243, %2244 : tensor<f32>\n",
      "    %2246 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2247 = mhlo.divide %2246, %2245 : tensor<f32>\n",
      "    %2248 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2249 = mhlo.multiply %2248, %2247 : tensor<f32>\n",
      "    %2250 = mhlo.cosine %2249 : tensor<f32>\n",
      "    %2251 = mhlo.sine %2249 : tensor<f32>\n",
      "    %2252 = mhlo.negate %2251 : tensor<f32>\n",
      "    %2253 = mhlo.convert %2250 : tensor<f32>\n",
      "    %2254 = mhlo.convert %2252 : tensor<f32>\n",
      "    %2255 = \"mhlo.broadcast_in_dim\"(%2253) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2256 = \"mhlo.broadcast_in_dim\"(%2254) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2257 = \"mhlo.concatenate\"(%2255, %2256) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2258 = mhlo.convert %2251 : tensor<f32>\n",
      "    %2259 = mhlo.convert %2250 : tensor<f32>\n",
      "    %2260 = \"mhlo.broadcast_in_dim\"(%2258) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2261 = \"mhlo.broadcast_in_dim\"(%2259) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2262 = \"mhlo.concatenate\"(%2260, %2261) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2263 = \"mhlo.broadcast_in_dim\"(%2257) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2264 = \"mhlo.broadcast_in_dim\"(%2262) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2265 = \"mhlo.concatenate\"(%2263, %2264) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2266 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2267 = mhlo.constant dense<7.500000e-01> : tensor<f32>\n",
      "    %2268 = mhlo.convert %2266 : tensor<f32>\n",
      "    %2269 = mhlo.convert %2267 : tensor<f32>\n",
      "    %2270 = mhlo.power %2268, %2269 : tensor<f32>\n",
      "    %2271 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2272 = mhlo.divide %2271, %2270 : tensor<f32>\n",
      "    %2273 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2274 = mhlo.multiply %2273, %2272 : tensor<f32>\n",
      "    %2275 = mhlo.cosine %2274 : tensor<f32>\n",
      "    %2276 = mhlo.sine %2274 : tensor<f32>\n",
      "    %2277 = mhlo.negate %2276 : tensor<f32>\n",
      "    %2278 = mhlo.convert %2275 : tensor<f32>\n",
      "    %2279 = mhlo.convert %2277 : tensor<f32>\n",
      "    %2280 = \"mhlo.broadcast_in_dim\"(%2278) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2281 = \"mhlo.broadcast_in_dim\"(%2279) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2282 = \"mhlo.concatenate\"(%2280, %2281) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2283 = mhlo.convert %2276 : tensor<f32>\n",
      "    %2284 = mhlo.convert %2275 : tensor<f32>\n",
      "    %2285 = \"mhlo.broadcast_in_dim\"(%2283) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2286 = \"mhlo.broadcast_in_dim\"(%2284) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2287 = \"mhlo.concatenate\"(%2285, %2286) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2288 = \"mhlo.broadcast_in_dim\"(%2282) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2289 = \"mhlo.broadcast_in_dim\"(%2287) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2290 = \"mhlo.concatenate\"(%2288, %2289) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2291 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2292 = mhlo.constant dense<0.791666686> : tensor<f32>\n",
      "    %2293 = mhlo.convert %2291 : tensor<f32>\n",
      "    %2294 = mhlo.convert %2292 : tensor<f32>\n",
      "    %2295 = mhlo.power %2293, %2294 : tensor<f32>\n",
      "    %2296 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2297 = mhlo.divide %2296, %2295 : tensor<f32>\n",
      "    %2298 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2299 = mhlo.multiply %2298, %2297 : tensor<f32>\n",
      "    %2300 = mhlo.cosine %2299 : tensor<f32>\n",
      "    %2301 = mhlo.sine %2299 : tensor<f32>\n",
      "    %2302 = mhlo.negate %2301 : tensor<f32>\n",
      "    %2303 = mhlo.convert %2300 : tensor<f32>\n",
      "    %2304 = mhlo.convert %2302 : tensor<f32>\n",
      "    %2305 = \"mhlo.broadcast_in_dim\"(%2303) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2306 = \"mhlo.broadcast_in_dim\"(%2304) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2307 = \"mhlo.concatenate\"(%2305, %2306) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2308 = mhlo.convert %2301 : tensor<f32>\n",
      "    %2309 = mhlo.convert %2300 : tensor<f32>\n",
      "    %2310 = \"mhlo.broadcast_in_dim\"(%2308) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2311 = \"mhlo.broadcast_in_dim\"(%2309) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2312 = \"mhlo.concatenate\"(%2310, %2311) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2313 = \"mhlo.broadcast_in_dim\"(%2307) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2314 = \"mhlo.broadcast_in_dim\"(%2312) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2315 = \"mhlo.concatenate\"(%2313, %2314) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2316 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2317 = mhlo.constant dense<0.833333313> : tensor<f32>\n",
      "    %2318 = mhlo.convert %2316 : tensor<f32>\n",
      "    %2319 = mhlo.convert %2317 : tensor<f32>\n",
      "    %2320 = mhlo.power %2318, %2319 : tensor<f32>\n",
      "    %2321 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2322 = mhlo.divide %2321, %2320 : tensor<f32>\n",
      "    %2323 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2324 = mhlo.multiply %2323, %2322 : tensor<f32>\n",
      "    %2325 = mhlo.cosine %2324 : tensor<f32>\n",
      "    %2326 = mhlo.sine %2324 : tensor<f32>\n",
      "    %2327 = mhlo.negate %2326 : tensor<f32>\n",
      "    %2328 = mhlo.convert %2325 : tensor<f32>\n",
      "    %2329 = mhlo.convert %2327 : tensor<f32>\n",
      "    %2330 = \"mhlo.broadcast_in_dim\"(%2328) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2331 = \"mhlo.broadcast_in_dim\"(%2329) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2332 = \"mhlo.concatenate\"(%2330, %2331) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2333 = mhlo.convert %2326 : tensor<f32>\n",
      "    %2334 = mhlo.convert %2325 : tensor<f32>\n",
      "    %2335 = \"mhlo.broadcast_in_dim\"(%2333) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2336 = \"mhlo.broadcast_in_dim\"(%2334) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2337 = \"mhlo.concatenate\"(%2335, %2336) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2338 = \"mhlo.broadcast_in_dim\"(%2332) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2339 = \"mhlo.broadcast_in_dim\"(%2337) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2340 = \"mhlo.concatenate\"(%2338, %2339) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2341 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2342 = mhlo.constant dense<8.750000e-01> : tensor<f32>\n",
      "    %2343 = mhlo.convert %2341 : tensor<f32>\n",
      "    %2344 = mhlo.convert %2342 : tensor<f32>\n",
      "    %2345 = mhlo.power %2343, %2344 : tensor<f32>\n",
      "    %2346 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2347 = mhlo.divide %2346, %2345 : tensor<f32>\n",
      "    %2348 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2349 = mhlo.multiply %2348, %2347 : tensor<f32>\n",
      "    %2350 = mhlo.cosine %2349 : tensor<f32>\n",
      "    %2351 = mhlo.sine %2349 : tensor<f32>\n",
      "    %2352 = mhlo.negate %2351 : tensor<f32>\n",
      "    %2353 = mhlo.convert %2350 : tensor<f32>\n",
      "    %2354 = mhlo.convert %2352 : tensor<f32>\n",
      "    %2355 = \"mhlo.broadcast_in_dim\"(%2353) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2356 = \"mhlo.broadcast_in_dim\"(%2354) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2357 = \"mhlo.concatenate\"(%2355, %2356) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2358 = mhlo.convert %2351 : tensor<f32>\n",
      "    %2359 = mhlo.convert %2350 : tensor<f32>\n",
      "    %2360 = \"mhlo.broadcast_in_dim\"(%2358) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2361 = \"mhlo.broadcast_in_dim\"(%2359) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2362 = \"mhlo.concatenate\"(%2360, %2361) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2363 = \"mhlo.broadcast_in_dim\"(%2357) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2364 = \"mhlo.broadcast_in_dim\"(%2362) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2365 = \"mhlo.concatenate\"(%2363, %2364) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2366 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2367 = mhlo.constant dense<0.916666686> : tensor<f32>\n",
      "    %2368 = mhlo.convert %2366 : tensor<f32>\n",
      "    %2369 = mhlo.convert %2367 : tensor<f32>\n",
      "    %2370 = mhlo.power %2368, %2369 : tensor<f32>\n",
      "    %2371 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2372 = mhlo.divide %2371, %2370 : tensor<f32>\n",
      "    %2373 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2374 = mhlo.multiply %2373, %2372 : tensor<f32>\n",
      "    %2375 = mhlo.cosine %2374 : tensor<f32>\n",
      "    %2376 = mhlo.sine %2374 : tensor<f32>\n",
      "    %2377 = mhlo.negate %2376 : tensor<f32>\n",
      "    %2378 = mhlo.convert %2375 : tensor<f32>\n",
      "    %2379 = mhlo.convert %2377 : tensor<f32>\n",
      "    %2380 = \"mhlo.broadcast_in_dim\"(%2378) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2381 = \"mhlo.broadcast_in_dim\"(%2379) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2382 = \"mhlo.concatenate\"(%2380, %2381) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2383 = mhlo.convert %2376 : tensor<f32>\n",
      "    %2384 = mhlo.convert %2375 : tensor<f32>\n",
      "    %2385 = \"mhlo.broadcast_in_dim\"(%2383) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2386 = \"mhlo.broadcast_in_dim\"(%2384) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2387 = \"mhlo.concatenate\"(%2385, %2386) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2388 = \"mhlo.broadcast_in_dim\"(%2382) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2389 = \"mhlo.broadcast_in_dim\"(%2387) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2390 = \"mhlo.concatenate\"(%2388, %2389) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2391 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2392 = mhlo.constant dense<0.958333313> : tensor<f32>\n",
      "    %2393 = mhlo.convert %2391 : tensor<f32>\n",
      "    %2394 = mhlo.convert %2392 : tensor<f32>\n",
      "    %2395 = mhlo.power %2393, %2394 : tensor<f32>\n",
      "    %2396 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2397 = mhlo.divide %2396, %2395 : tensor<f32>\n",
      "    %2398 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2399 = mhlo.multiply %2398, %2397 : tensor<f32>\n",
      "    %2400 = mhlo.cosine %2399 : tensor<f32>\n",
      "    %2401 = mhlo.sine %2399 : tensor<f32>\n",
      "    %2402 = mhlo.negate %2401 : tensor<f32>\n",
      "    %2403 = mhlo.convert %2400 : tensor<f32>\n",
      "    %2404 = mhlo.convert %2402 : tensor<f32>\n",
      "    %2405 = \"mhlo.broadcast_in_dim\"(%2403) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2406 = \"mhlo.broadcast_in_dim\"(%2404) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2407 = \"mhlo.concatenate\"(%2405, %2406) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2408 = mhlo.convert %2401 : tensor<f32>\n",
      "    %2409 = mhlo.convert %2400 : tensor<f32>\n",
      "    %2410 = \"mhlo.broadcast_in_dim\"(%2408) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2411 = \"mhlo.broadcast_in_dim\"(%2409) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2412 = \"mhlo.concatenate\"(%2410, %2411) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2413 = \"mhlo.broadcast_in_dim\"(%2407) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2414 = \"mhlo.broadcast_in_dim\"(%2412) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2415 = \"mhlo.concatenate\"(%2413, %2414) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2416 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2417 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2418 = mhlo.convert %2416 : tensor<f32>\n",
      "    %2419 = mhlo.convert %2417 : tensor<f32>\n",
      "    %2420 = mhlo.power %2418, %2419 : tensor<f32>\n",
      "    %2421 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2422 = mhlo.divide %2421, %2420 : tensor<f32>\n",
      "    %2423 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2424 = mhlo.multiply %2423, %2422 : tensor<f32>\n",
      "    %2425 = mhlo.cosine %2424 : tensor<f32>\n",
      "    %2426 = mhlo.sine %2424 : tensor<f32>\n",
      "    %2427 = mhlo.negate %2426 : tensor<f32>\n",
      "    %2428 = mhlo.convert %2425 : tensor<f32>\n",
      "    %2429 = mhlo.convert %2427 : tensor<f32>\n",
      "    %2430 = \"mhlo.broadcast_in_dim\"(%2428) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2431 = \"mhlo.broadcast_in_dim\"(%2429) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2432 = \"mhlo.concatenate\"(%2430, %2431) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2433 = mhlo.convert %2426 : tensor<f32>\n",
      "    %2434 = mhlo.convert %2425 : tensor<f32>\n",
      "    %2435 = \"mhlo.broadcast_in_dim\"(%2433) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2436 = \"mhlo.broadcast_in_dim\"(%2434) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2437 = \"mhlo.concatenate\"(%2435, %2436) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2438 = \"mhlo.broadcast_in_dim\"(%2432) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2439 = \"mhlo.broadcast_in_dim\"(%2437) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2440 = \"mhlo.concatenate\"(%2438, %2439) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2441 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2442 = mhlo.constant dense<0.0416666679> : tensor<f32>\n",
      "    %2443 = mhlo.convert %2441 : tensor<f32>\n",
      "    %2444 = mhlo.convert %2442 : tensor<f32>\n",
      "    %2445 = mhlo.power %2443, %2444 : tensor<f32>\n",
      "    %2446 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2447 = mhlo.divide %2446, %2445 : tensor<f32>\n",
      "    %2448 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2449 = mhlo.multiply %2448, %2447 : tensor<f32>\n",
      "    %2450 = mhlo.cosine %2449 : tensor<f32>\n",
      "    %2451 = mhlo.sine %2449 : tensor<f32>\n",
      "    %2452 = mhlo.negate %2451 : tensor<f32>\n",
      "    %2453 = mhlo.convert %2450 : tensor<f32>\n",
      "    %2454 = mhlo.convert %2452 : tensor<f32>\n",
      "    %2455 = \"mhlo.broadcast_in_dim\"(%2453) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2456 = \"mhlo.broadcast_in_dim\"(%2454) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2457 = \"mhlo.concatenate\"(%2455, %2456) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2458 = mhlo.convert %2451 : tensor<f32>\n",
      "    %2459 = mhlo.convert %2450 : tensor<f32>\n",
      "    %2460 = \"mhlo.broadcast_in_dim\"(%2458) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2461 = \"mhlo.broadcast_in_dim\"(%2459) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2462 = \"mhlo.concatenate\"(%2460, %2461) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2463 = \"mhlo.broadcast_in_dim\"(%2457) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2464 = \"mhlo.broadcast_in_dim\"(%2462) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2465 = \"mhlo.concatenate\"(%2463, %2464) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2466 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2467 = mhlo.constant dense<0.0833333358> : tensor<f32>\n",
      "    %2468 = mhlo.convert %2466 : tensor<f32>\n",
      "    %2469 = mhlo.convert %2467 : tensor<f32>\n",
      "    %2470 = mhlo.power %2468, %2469 : tensor<f32>\n",
      "    %2471 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2472 = mhlo.divide %2471, %2470 : tensor<f32>\n",
      "    %2473 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2474 = mhlo.multiply %2473, %2472 : tensor<f32>\n",
      "    %2475 = mhlo.cosine %2474 : tensor<f32>\n",
      "    %2476 = mhlo.sine %2474 : tensor<f32>\n",
      "    %2477 = mhlo.negate %2476 : tensor<f32>\n",
      "    %2478 = mhlo.convert %2475 : tensor<f32>\n",
      "    %2479 = mhlo.convert %2477 : tensor<f32>\n",
      "    %2480 = \"mhlo.broadcast_in_dim\"(%2478) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2481 = \"mhlo.broadcast_in_dim\"(%2479) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2482 = \"mhlo.concatenate\"(%2480, %2481) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2483 = mhlo.convert %2476 : tensor<f32>\n",
      "    %2484 = mhlo.convert %2475 : tensor<f32>\n",
      "    %2485 = \"mhlo.broadcast_in_dim\"(%2483) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2486 = \"mhlo.broadcast_in_dim\"(%2484) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2487 = \"mhlo.concatenate\"(%2485, %2486) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2488 = \"mhlo.broadcast_in_dim\"(%2482) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2489 = \"mhlo.broadcast_in_dim\"(%2487) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2490 = \"mhlo.concatenate\"(%2488, %2489) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2491 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2492 = mhlo.constant dense<1.250000e-01> : tensor<f32>\n",
      "    %2493 = mhlo.convert %2491 : tensor<f32>\n",
      "    %2494 = mhlo.convert %2492 : tensor<f32>\n",
      "    %2495 = mhlo.power %2493, %2494 : tensor<f32>\n",
      "    %2496 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2497 = mhlo.divide %2496, %2495 : tensor<f32>\n",
      "    %2498 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2499 = mhlo.multiply %2498, %2497 : tensor<f32>\n",
      "    %2500 = mhlo.cosine %2499 : tensor<f32>\n",
      "    %2501 = mhlo.sine %2499 : tensor<f32>\n",
      "    %2502 = mhlo.negate %2501 : tensor<f32>\n",
      "    %2503 = mhlo.convert %2500 : tensor<f32>\n",
      "    %2504 = mhlo.convert %2502 : tensor<f32>\n",
      "    %2505 = \"mhlo.broadcast_in_dim\"(%2503) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2506 = \"mhlo.broadcast_in_dim\"(%2504) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2507 = \"mhlo.concatenate\"(%2505, %2506) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2508 = mhlo.convert %2501 : tensor<f32>\n",
      "    %2509 = mhlo.convert %2500 : tensor<f32>\n",
      "    %2510 = \"mhlo.broadcast_in_dim\"(%2508) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2511 = \"mhlo.broadcast_in_dim\"(%2509) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2512 = \"mhlo.concatenate\"(%2510, %2511) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2513 = \"mhlo.broadcast_in_dim\"(%2507) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2514 = \"mhlo.broadcast_in_dim\"(%2512) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2515 = \"mhlo.concatenate\"(%2513, %2514) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2516 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2517 = mhlo.constant dense<0.166666672> : tensor<f32>\n",
      "    %2518 = mhlo.convert %2516 : tensor<f32>\n",
      "    %2519 = mhlo.convert %2517 : tensor<f32>\n",
      "    %2520 = mhlo.power %2518, %2519 : tensor<f32>\n",
      "    %2521 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2522 = mhlo.divide %2521, %2520 : tensor<f32>\n",
      "    %2523 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2524 = mhlo.multiply %2523, %2522 : tensor<f32>\n",
      "    %2525 = mhlo.cosine %2524 : tensor<f32>\n",
      "    %2526 = mhlo.sine %2524 : tensor<f32>\n",
      "    %2527 = mhlo.negate %2526 : tensor<f32>\n",
      "    %2528 = mhlo.convert %2525 : tensor<f32>\n",
      "    %2529 = mhlo.convert %2527 : tensor<f32>\n",
      "    %2530 = \"mhlo.broadcast_in_dim\"(%2528) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2531 = \"mhlo.broadcast_in_dim\"(%2529) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2532 = \"mhlo.concatenate\"(%2530, %2531) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2533 = mhlo.convert %2526 : tensor<f32>\n",
      "    %2534 = mhlo.convert %2525 : tensor<f32>\n",
      "    %2535 = \"mhlo.broadcast_in_dim\"(%2533) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2536 = \"mhlo.broadcast_in_dim\"(%2534) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2537 = \"mhlo.concatenate\"(%2535, %2536) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2538 = \"mhlo.broadcast_in_dim\"(%2532) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2539 = \"mhlo.broadcast_in_dim\"(%2537) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2540 = \"mhlo.concatenate\"(%2538, %2539) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2541 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2542 = mhlo.constant dense<0.208333328> : tensor<f32>\n",
      "    %2543 = mhlo.convert %2541 : tensor<f32>\n",
      "    %2544 = mhlo.convert %2542 : tensor<f32>\n",
      "    %2545 = mhlo.power %2543, %2544 : tensor<f32>\n",
      "    %2546 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2547 = mhlo.divide %2546, %2545 : tensor<f32>\n",
      "    %2548 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2549 = mhlo.multiply %2548, %2547 : tensor<f32>\n",
      "    %2550 = mhlo.cosine %2549 : tensor<f32>\n",
      "    %2551 = mhlo.sine %2549 : tensor<f32>\n",
      "    %2552 = mhlo.negate %2551 : tensor<f32>\n",
      "    %2553 = mhlo.convert %2550 : tensor<f32>\n",
      "    %2554 = mhlo.convert %2552 : tensor<f32>\n",
      "    %2555 = \"mhlo.broadcast_in_dim\"(%2553) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2556 = \"mhlo.broadcast_in_dim\"(%2554) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2557 = \"mhlo.concatenate\"(%2555, %2556) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2558 = mhlo.convert %2551 : tensor<f32>\n",
      "    %2559 = mhlo.convert %2550 : tensor<f32>\n",
      "    %2560 = \"mhlo.broadcast_in_dim\"(%2558) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2561 = \"mhlo.broadcast_in_dim\"(%2559) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2562 = \"mhlo.concatenate\"(%2560, %2561) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2563 = \"mhlo.broadcast_in_dim\"(%2557) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2564 = \"mhlo.broadcast_in_dim\"(%2562) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2565 = \"mhlo.concatenate\"(%2563, %2564) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2566 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2567 = mhlo.constant dense<2.500000e-01> : tensor<f32>\n",
      "    %2568 = mhlo.convert %2566 : tensor<f32>\n",
      "    %2569 = mhlo.convert %2567 : tensor<f32>\n",
      "    %2570 = mhlo.power %2568, %2569 : tensor<f32>\n",
      "    %2571 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2572 = mhlo.divide %2571, %2570 : tensor<f32>\n",
      "    %2573 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2574 = mhlo.multiply %2573, %2572 : tensor<f32>\n",
      "    %2575 = mhlo.cosine %2574 : tensor<f32>\n",
      "    %2576 = mhlo.sine %2574 : tensor<f32>\n",
      "    %2577 = mhlo.negate %2576 : tensor<f32>\n",
      "    %2578 = mhlo.convert %2575 : tensor<f32>\n",
      "    %2579 = mhlo.convert %2577 : tensor<f32>\n",
      "    %2580 = \"mhlo.broadcast_in_dim\"(%2578) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2581 = \"mhlo.broadcast_in_dim\"(%2579) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2582 = \"mhlo.concatenate\"(%2580, %2581) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2583 = mhlo.convert %2576 : tensor<f32>\n",
      "    %2584 = mhlo.convert %2575 : tensor<f32>\n",
      "    %2585 = \"mhlo.broadcast_in_dim\"(%2583) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2586 = \"mhlo.broadcast_in_dim\"(%2584) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2587 = \"mhlo.concatenate\"(%2585, %2586) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2588 = \"mhlo.broadcast_in_dim\"(%2582) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2589 = \"mhlo.broadcast_in_dim\"(%2587) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2590 = \"mhlo.concatenate\"(%2588, %2589) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2591 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2592 = mhlo.constant dense<0.291666657> : tensor<f32>\n",
      "    %2593 = mhlo.convert %2591 : tensor<f32>\n",
      "    %2594 = mhlo.convert %2592 : tensor<f32>\n",
      "    %2595 = mhlo.power %2593, %2594 : tensor<f32>\n",
      "    %2596 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2597 = mhlo.divide %2596, %2595 : tensor<f32>\n",
      "    %2598 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2599 = mhlo.multiply %2598, %2597 : tensor<f32>\n",
      "    %2600 = mhlo.cosine %2599 : tensor<f32>\n",
      "    %2601 = mhlo.sine %2599 : tensor<f32>\n",
      "    %2602 = mhlo.negate %2601 : tensor<f32>\n",
      "    %2603 = mhlo.convert %2600 : tensor<f32>\n",
      "    %2604 = mhlo.convert %2602 : tensor<f32>\n",
      "    %2605 = \"mhlo.broadcast_in_dim\"(%2603) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2606 = \"mhlo.broadcast_in_dim\"(%2604) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2607 = \"mhlo.concatenate\"(%2605, %2606) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2608 = mhlo.convert %2601 : tensor<f32>\n",
      "    %2609 = mhlo.convert %2600 : tensor<f32>\n",
      "    %2610 = \"mhlo.broadcast_in_dim\"(%2608) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2611 = \"mhlo.broadcast_in_dim\"(%2609) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2612 = \"mhlo.concatenate\"(%2610, %2611) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2613 = \"mhlo.broadcast_in_dim\"(%2607) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2614 = \"mhlo.broadcast_in_dim\"(%2612) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2615 = \"mhlo.concatenate\"(%2613, %2614) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2616 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2617 = mhlo.constant dense<0.333333343> : tensor<f32>\n",
      "    %2618 = mhlo.convert %2616 : tensor<f32>\n",
      "    %2619 = mhlo.convert %2617 : tensor<f32>\n",
      "    %2620 = mhlo.power %2618, %2619 : tensor<f32>\n",
      "    %2621 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2622 = mhlo.divide %2621, %2620 : tensor<f32>\n",
      "    %2623 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2624 = mhlo.multiply %2623, %2622 : tensor<f32>\n",
      "    %2625 = mhlo.cosine %2624 : tensor<f32>\n",
      "    %2626 = mhlo.sine %2624 : tensor<f32>\n",
      "    %2627 = mhlo.negate %2626 : tensor<f32>\n",
      "    %2628 = mhlo.convert %2625 : tensor<f32>\n",
      "    %2629 = mhlo.convert %2627 : tensor<f32>\n",
      "    %2630 = \"mhlo.broadcast_in_dim\"(%2628) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2631 = \"mhlo.broadcast_in_dim\"(%2629) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2632 = \"mhlo.concatenate\"(%2630, %2631) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2633 = mhlo.convert %2626 : tensor<f32>\n",
      "    %2634 = mhlo.convert %2625 : tensor<f32>\n",
      "    %2635 = \"mhlo.broadcast_in_dim\"(%2633) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2636 = \"mhlo.broadcast_in_dim\"(%2634) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2637 = \"mhlo.concatenate\"(%2635, %2636) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2638 = \"mhlo.broadcast_in_dim\"(%2632) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2639 = \"mhlo.broadcast_in_dim\"(%2637) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2640 = \"mhlo.concatenate\"(%2638, %2639) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2641 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2642 = mhlo.constant dense<3.750000e-01> : tensor<f32>\n",
      "    %2643 = mhlo.convert %2641 : tensor<f32>\n",
      "    %2644 = mhlo.convert %2642 : tensor<f32>\n",
      "    %2645 = mhlo.power %2643, %2644 : tensor<f32>\n",
      "    %2646 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2647 = mhlo.divide %2646, %2645 : tensor<f32>\n",
      "    %2648 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2649 = mhlo.multiply %2648, %2647 : tensor<f32>\n",
      "    %2650 = mhlo.cosine %2649 : tensor<f32>\n",
      "    %2651 = mhlo.sine %2649 : tensor<f32>\n",
      "    %2652 = mhlo.negate %2651 : tensor<f32>\n",
      "    %2653 = mhlo.convert %2650 : tensor<f32>\n",
      "    %2654 = mhlo.convert %2652 : tensor<f32>\n",
      "    %2655 = \"mhlo.broadcast_in_dim\"(%2653) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2656 = \"mhlo.broadcast_in_dim\"(%2654) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2657 = \"mhlo.concatenate\"(%2655, %2656) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2658 = mhlo.convert %2651 : tensor<f32>\n",
      "    %2659 = mhlo.convert %2650 : tensor<f32>\n",
      "    %2660 = \"mhlo.broadcast_in_dim\"(%2658) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2661 = \"mhlo.broadcast_in_dim\"(%2659) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2662 = \"mhlo.concatenate\"(%2660, %2661) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2663 = \"mhlo.broadcast_in_dim\"(%2657) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2664 = \"mhlo.broadcast_in_dim\"(%2662) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2665 = \"mhlo.concatenate\"(%2663, %2664) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2666 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2667 = mhlo.constant dense<0.416666657> : tensor<f32>\n",
      "    %2668 = mhlo.convert %2666 : tensor<f32>\n",
      "    %2669 = mhlo.convert %2667 : tensor<f32>\n",
      "    %2670 = mhlo.power %2668, %2669 : tensor<f32>\n",
      "    %2671 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2672 = mhlo.divide %2671, %2670 : tensor<f32>\n",
      "    %2673 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2674 = mhlo.multiply %2673, %2672 : tensor<f32>\n",
      "    %2675 = mhlo.cosine %2674 : tensor<f32>\n",
      "    %2676 = mhlo.sine %2674 : tensor<f32>\n",
      "    %2677 = mhlo.negate %2676 : tensor<f32>\n",
      "    %2678 = mhlo.convert %2675 : tensor<f32>\n",
      "    %2679 = mhlo.convert %2677 : tensor<f32>\n",
      "    %2680 = \"mhlo.broadcast_in_dim\"(%2678) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2681 = \"mhlo.broadcast_in_dim\"(%2679) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2682 = \"mhlo.concatenate\"(%2680, %2681) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2683 = mhlo.convert %2676 : tensor<f32>\n",
      "    %2684 = mhlo.convert %2675 : tensor<f32>\n",
      "    %2685 = \"mhlo.broadcast_in_dim\"(%2683) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2686 = \"mhlo.broadcast_in_dim\"(%2684) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2687 = \"mhlo.concatenate\"(%2685, %2686) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2688 = \"mhlo.broadcast_in_dim\"(%2682) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2689 = \"mhlo.broadcast_in_dim\"(%2687) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2690 = \"mhlo.concatenate\"(%2688, %2689) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2691 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2692 = mhlo.constant dense<0.458333343> : tensor<f32>\n",
      "    %2693 = mhlo.convert %2691 : tensor<f32>\n",
      "    %2694 = mhlo.convert %2692 : tensor<f32>\n",
      "    %2695 = mhlo.power %2693, %2694 : tensor<f32>\n",
      "    %2696 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2697 = mhlo.divide %2696, %2695 : tensor<f32>\n",
      "    %2698 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2699 = mhlo.multiply %2698, %2697 : tensor<f32>\n",
      "    %2700 = mhlo.cosine %2699 : tensor<f32>\n",
      "    %2701 = mhlo.sine %2699 : tensor<f32>\n",
      "    %2702 = mhlo.negate %2701 : tensor<f32>\n",
      "    %2703 = mhlo.convert %2700 : tensor<f32>\n",
      "    %2704 = mhlo.convert %2702 : tensor<f32>\n",
      "    %2705 = \"mhlo.broadcast_in_dim\"(%2703) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2706 = \"mhlo.broadcast_in_dim\"(%2704) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2707 = \"mhlo.concatenate\"(%2705, %2706) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2708 = mhlo.convert %2701 : tensor<f32>\n",
      "    %2709 = mhlo.convert %2700 : tensor<f32>\n",
      "    %2710 = \"mhlo.broadcast_in_dim\"(%2708) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2711 = \"mhlo.broadcast_in_dim\"(%2709) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2712 = \"mhlo.concatenate\"(%2710, %2711) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2713 = \"mhlo.broadcast_in_dim\"(%2707) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2714 = \"mhlo.broadcast_in_dim\"(%2712) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2715 = \"mhlo.concatenate\"(%2713, %2714) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2716 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2717 = mhlo.constant dense<5.000000e-01> : tensor<f32>\n",
      "    %2718 = mhlo.convert %2716 : tensor<f32>\n",
      "    %2719 = mhlo.convert %2717 : tensor<f32>\n",
      "    %2720 = mhlo.power %2718, %2719 : tensor<f32>\n",
      "    %2721 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2722 = mhlo.divide %2721, %2720 : tensor<f32>\n",
      "    %2723 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2724 = mhlo.multiply %2723, %2722 : tensor<f32>\n",
      "    %2725 = mhlo.cosine %2724 : tensor<f32>\n",
      "    %2726 = mhlo.sine %2724 : tensor<f32>\n",
      "    %2727 = mhlo.negate %2726 : tensor<f32>\n",
      "    %2728 = mhlo.convert %2725 : tensor<f32>\n",
      "    %2729 = mhlo.convert %2727 : tensor<f32>\n",
      "    %2730 = \"mhlo.broadcast_in_dim\"(%2728) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2731 = \"mhlo.broadcast_in_dim\"(%2729) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2732 = \"mhlo.concatenate\"(%2730, %2731) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2733 = mhlo.convert %2726 : tensor<f32>\n",
      "    %2734 = mhlo.convert %2725 : tensor<f32>\n",
      "    %2735 = \"mhlo.broadcast_in_dim\"(%2733) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2736 = \"mhlo.broadcast_in_dim\"(%2734) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2737 = \"mhlo.concatenate\"(%2735, %2736) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2738 = \"mhlo.broadcast_in_dim\"(%2732) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2739 = \"mhlo.broadcast_in_dim\"(%2737) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2740 = \"mhlo.concatenate\"(%2738, %2739) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2741 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2742 = mhlo.constant dense<0.541666687> : tensor<f32>\n",
      "    %2743 = mhlo.convert %2741 : tensor<f32>\n",
      "    %2744 = mhlo.convert %2742 : tensor<f32>\n",
      "    %2745 = mhlo.power %2743, %2744 : tensor<f32>\n",
      "    %2746 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2747 = mhlo.divide %2746, %2745 : tensor<f32>\n",
      "    %2748 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2749 = mhlo.multiply %2748, %2747 : tensor<f32>\n",
      "    %2750 = mhlo.cosine %2749 : tensor<f32>\n",
      "    %2751 = mhlo.sine %2749 : tensor<f32>\n",
      "    %2752 = mhlo.negate %2751 : tensor<f32>\n",
      "    %2753 = mhlo.convert %2750 : tensor<f32>\n",
      "    %2754 = mhlo.convert %2752 : tensor<f32>\n",
      "    %2755 = \"mhlo.broadcast_in_dim\"(%2753) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2756 = \"mhlo.broadcast_in_dim\"(%2754) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2757 = \"mhlo.concatenate\"(%2755, %2756) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2758 = mhlo.convert %2751 : tensor<f32>\n",
      "    %2759 = mhlo.convert %2750 : tensor<f32>\n",
      "    %2760 = \"mhlo.broadcast_in_dim\"(%2758) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2761 = \"mhlo.broadcast_in_dim\"(%2759) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2762 = \"mhlo.concatenate\"(%2760, %2761) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2763 = \"mhlo.broadcast_in_dim\"(%2757) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2764 = \"mhlo.broadcast_in_dim\"(%2762) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2765 = \"mhlo.concatenate\"(%2763, %2764) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2766 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2767 = mhlo.constant dense<0.583333313> : tensor<f32>\n",
      "    %2768 = mhlo.convert %2766 : tensor<f32>\n",
      "    %2769 = mhlo.convert %2767 : tensor<f32>\n",
      "    %2770 = mhlo.power %2768, %2769 : tensor<f32>\n",
      "    %2771 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2772 = mhlo.divide %2771, %2770 : tensor<f32>\n",
      "    %2773 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2774 = mhlo.multiply %2773, %2772 : tensor<f32>\n",
      "    %2775 = mhlo.cosine %2774 : tensor<f32>\n",
      "    %2776 = mhlo.sine %2774 : tensor<f32>\n",
      "    %2777 = mhlo.negate %2776 : tensor<f32>\n",
      "    %2778 = mhlo.convert %2775 : tensor<f32>\n",
      "    %2779 = mhlo.convert %2777 : tensor<f32>\n",
      "    %2780 = \"mhlo.broadcast_in_dim\"(%2778) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2781 = \"mhlo.broadcast_in_dim\"(%2779) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2782 = \"mhlo.concatenate\"(%2780, %2781) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2783 = mhlo.convert %2776 : tensor<f32>\n",
      "    %2784 = mhlo.convert %2775 : tensor<f32>\n",
      "    %2785 = \"mhlo.broadcast_in_dim\"(%2783) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2786 = \"mhlo.broadcast_in_dim\"(%2784) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2787 = \"mhlo.concatenate\"(%2785, %2786) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2788 = \"mhlo.broadcast_in_dim\"(%2782) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2789 = \"mhlo.broadcast_in_dim\"(%2787) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2790 = \"mhlo.concatenate\"(%2788, %2789) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2791 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2792 = mhlo.constant dense<6.250000e-01> : tensor<f32>\n",
      "    %2793 = mhlo.convert %2791 : tensor<f32>\n",
      "    %2794 = mhlo.convert %2792 : tensor<f32>\n",
      "    %2795 = mhlo.power %2793, %2794 : tensor<f32>\n",
      "    %2796 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2797 = mhlo.divide %2796, %2795 : tensor<f32>\n",
      "    %2798 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2799 = mhlo.multiply %2798, %2797 : tensor<f32>\n",
      "    %2800 = mhlo.cosine %2799 : tensor<f32>\n",
      "    %2801 = mhlo.sine %2799 : tensor<f32>\n",
      "    %2802 = mhlo.negate %2801 : tensor<f32>\n",
      "    %2803 = mhlo.convert %2800 : tensor<f32>\n",
      "    %2804 = mhlo.convert %2802 : tensor<f32>\n",
      "    %2805 = \"mhlo.broadcast_in_dim\"(%2803) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2806 = \"mhlo.broadcast_in_dim\"(%2804) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2807 = \"mhlo.concatenate\"(%2805, %2806) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2808 = mhlo.convert %2801 : tensor<f32>\n",
      "    %2809 = mhlo.convert %2800 : tensor<f32>\n",
      "    %2810 = \"mhlo.broadcast_in_dim\"(%2808) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2811 = \"mhlo.broadcast_in_dim\"(%2809) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2812 = \"mhlo.concatenate\"(%2810, %2811) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2813 = \"mhlo.broadcast_in_dim\"(%2807) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2814 = \"mhlo.broadcast_in_dim\"(%2812) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2815 = \"mhlo.concatenate\"(%2813, %2814) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2816 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2817 = mhlo.constant dense<0.666666686> : tensor<f32>\n",
      "    %2818 = mhlo.convert %2816 : tensor<f32>\n",
      "    %2819 = mhlo.convert %2817 : tensor<f32>\n",
      "    %2820 = mhlo.power %2818, %2819 : tensor<f32>\n",
      "    %2821 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2822 = mhlo.divide %2821, %2820 : tensor<f32>\n",
      "    %2823 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2824 = mhlo.multiply %2823, %2822 : tensor<f32>\n",
      "    %2825 = mhlo.cosine %2824 : tensor<f32>\n",
      "    %2826 = mhlo.sine %2824 : tensor<f32>\n",
      "    %2827 = mhlo.negate %2826 : tensor<f32>\n",
      "    %2828 = mhlo.convert %2825 : tensor<f32>\n",
      "    %2829 = mhlo.convert %2827 : tensor<f32>\n",
      "    %2830 = \"mhlo.broadcast_in_dim\"(%2828) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2831 = \"mhlo.broadcast_in_dim\"(%2829) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2832 = \"mhlo.concatenate\"(%2830, %2831) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2833 = mhlo.convert %2826 : tensor<f32>\n",
      "    %2834 = mhlo.convert %2825 : tensor<f32>\n",
      "    %2835 = \"mhlo.broadcast_in_dim\"(%2833) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2836 = \"mhlo.broadcast_in_dim\"(%2834) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2837 = \"mhlo.concatenate\"(%2835, %2836) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2838 = \"mhlo.broadcast_in_dim\"(%2832) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2839 = \"mhlo.broadcast_in_dim\"(%2837) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2840 = \"mhlo.concatenate\"(%2838, %2839) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2841 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2842 = mhlo.constant dense<0.708333313> : tensor<f32>\n",
      "    %2843 = mhlo.convert %2841 : tensor<f32>\n",
      "    %2844 = mhlo.convert %2842 : tensor<f32>\n",
      "    %2845 = mhlo.power %2843, %2844 : tensor<f32>\n",
      "    %2846 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2847 = mhlo.divide %2846, %2845 : tensor<f32>\n",
      "    %2848 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2849 = mhlo.multiply %2848, %2847 : tensor<f32>\n",
      "    %2850 = mhlo.cosine %2849 : tensor<f32>\n",
      "    %2851 = mhlo.sine %2849 : tensor<f32>\n",
      "    %2852 = mhlo.negate %2851 : tensor<f32>\n",
      "    %2853 = mhlo.convert %2850 : tensor<f32>\n",
      "    %2854 = mhlo.convert %2852 : tensor<f32>\n",
      "    %2855 = \"mhlo.broadcast_in_dim\"(%2853) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2856 = \"mhlo.broadcast_in_dim\"(%2854) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2857 = \"mhlo.concatenate\"(%2855, %2856) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2858 = mhlo.convert %2851 : tensor<f32>\n",
      "    %2859 = mhlo.convert %2850 : tensor<f32>\n",
      "    %2860 = \"mhlo.broadcast_in_dim\"(%2858) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2861 = \"mhlo.broadcast_in_dim\"(%2859) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2862 = \"mhlo.concatenate\"(%2860, %2861) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2863 = \"mhlo.broadcast_in_dim\"(%2857) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2864 = \"mhlo.broadcast_in_dim\"(%2862) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2865 = \"mhlo.concatenate\"(%2863, %2864) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2866 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2867 = mhlo.constant dense<7.500000e-01> : tensor<f32>\n",
      "    %2868 = mhlo.convert %2866 : tensor<f32>\n",
      "    %2869 = mhlo.convert %2867 : tensor<f32>\n",
      "    %2870 = mhlo.power %2868, %2869 : tensor<f32>\n",
      "    %2871 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2872 = mhlo.divide %2871, %2870 : tensor<f32>\n",
      "    %2873 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2874 = mhlo.multiply %2873, %2872 : tensor<f32>\n",
      "    %2875 = mhlo.cosine %2874 : tensor<f32>\n",
      "    %2876 = mhlo.sine %2874 : tensor<f32>\n",
      "    %2877 = mhlo.negate %2876 : tensor<f32>\n",
      "    %2878 = mhlo.convert %2875 : tensor<f32>\n",
      "    %2879 = mhlo.convert %2877 : tensor<f32>\n",
      "    %2880 = \"mhlo.broadcast_in_dim\"(%2878) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2881 = \"mhlo.broadcast_in_dim\"(%2879) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2882 = \"mhlo.concatenate\"(%2880, %2881) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2883 = mhlo.convert %2876 : tensor<f32>\n",
      "    %2884 = mhlo.convert %2875 : tensor<f32>\n",
      "    %2885 = \"mhlo.broadcast_in_dim\"(%2883) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2886 = \"mhlo.broadcast_in_dim\"(%2884) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2887 = \"mhlo.concatenate\"(%2885, %2886) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2888 = \"mhlo.broadcast_in_dim\"(%2882) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2889 = \"mhlo.broadcast_in_dim\"(%2887) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2890 = \"mhlo.concatenate\"(%2888, %2889) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2891 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2892 = mhlo.constant dense<0.791666686> : tensor<f32>\n",
      "    %2893 = mhlo.convert %2891 : tensor<f32>\n",
      "    %2894 = mhlo.convert %2892 : tensor<f32>\n",
      "    %2895 = mhlo.power %2893, %2894 : tensor<f32>\n",
      "    %2896 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2897 = mhlo.divide %2896, %2895 : tensor<f32>\n",
      "    %2898 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2899 = mhlo.multiply %2898, %2897 : tensor<f32>\n",
      "    %2900 = mhlo.cosine %2899 : tensor<f32>\n",
      "    %2901 = mhlo.sine %2899 : tensor<f32>\n",
      "    %2902 = mhlo.negate %2901 : tensor<f32>\n",
      "    %2903 = mhlo.convert %2900 : tensor<f32>\n",
      "    %2904 = mhlo.convert %2902 : tensor<f32>\n",
      "    %2905 = \"mhlo.broadcast_in_dim\"(%2903) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2906 = \"mhlo.broadcast_in_dim\"(%2904) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2907 = \"mhlo.concatenate\"(%2905, %2906) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2908 = mhlo.convert %2901 : tensor<f32>\n",
      "    %2909 = mhlo.convert %2900 : tensor<f32>\n",
      "    %2910 = \"mhlo.broadcast_in_dim\"(%2908) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2911 = \"mhlo.broadcast_in_dim\"(%2909) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2912 = \"mhlo.concatenate\"(%2910, %2911) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2913 = \"mhlo.broadcast_in_dim\"(%2907) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2914 = \"mhlo.broadcast_in_dim\"(%2912) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2915 = \"mhlo.concatenate\"(%2913, %2914) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2916 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2917 = mhlo.constant dense<0.833333313> : tensor<f32>\n",
      "    %2918 = mhlo.convert %2916 : tensor<f32>\n",
      "    %2919 = mhlo.convert %2917 : tensor<f32>\n",
      "    %2920 = mhlo.power %2918, %2919 : tensor<f32>\n",
      "    %2921 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2922 = mhlo.divide %2921, %2920 : tensor<f32>\n",
      "    %2923 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2924 = mhlo.multiply %2923, %2922 : tensor<f32>\n",
      "    %2925 = mhlo.cosine %2924 : tensor<f32>\n",
      "    %2926 = mhlo.sine %2924 : tensor<f32>\n",
      "    %2927 = mhlo.negate %2926 : tensor<f32>\n",
      "    %2928 = mhlo.convert %2925 : tensor<f32>\n",
      "    %2929 = mhlo.convert %2927 : tensor<f32>\n",
      "    %2930 = \"mhlo.broadcast_in_dim\"(%2928) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2931 = \"mhlo.broadcast_in_dim\"(%2929) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2932 = \"mhlo.concatenate\"(%2930, %2931) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2933 = mhlo.convert %2926 : tensor<f32>\n",
      "    %2934 = mhlo.convert %2925 : tensor<f32>\n",
      "    %2935 = \"mhlo.broadcast_in_dim\"(%2933) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2936 = \"mhlo.broadcast_in_dim\"(%2934) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2937 = \"mhlo.concatenate\"(%2935, %2936) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2938 = \"mhlo.broadcast_in_dim\"(%2932) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2939 = \"mhlo.broadcast_in_dim\"(%2937) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2940 = \"mhlo.concatenate\"(%2938, %2939) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2941 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2942 = mhlo.constant dense<8.750000e-01> : tensor<f32>\n",
      "    %2943 = mhlo.convert %2941 : tensor<f32>\n",
      "    %2944 = mhlo.convert %2942 : tensor<f32>\n",
      "    %2945 = mhlo.power %2943, %2944 : tensor<f32>\n",
      "    %2946 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2947 = mhlo.divide %2946, %2945 : tensor<f32>\n",
      "    %2948 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2949 = mhlo.multiply %2948, %2947 : tensor<f32>\n",
      "    %2950 = mhlo.cosine %2949 : tensor<f32>\n",
      "    %2951 = mhlo.sine %2949 : tensor<f32>\n",
      "    %2952 = mhlo.negate %2951 : tensor<f32>\n",
      "    %2953 = mhlo.convert %2950 : tensor<f32>\n",
      "    %2954 = mhlo.convert %2952 : tensor<f32>\n",
      "    %2955 = \"mhlo.broadcast_in_dim\"(%2953) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2956 = \"mhlo.broadcast_in_dim\"(%2954) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2957 = \"mhlo.concatenate\"(%2955, %2956) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2958 = mhlo.convert %2951 : tensor<f32>\n",
      "    %2959 = mhlo.convert %2950 : tensor<f32>\n",
      "    %2960 = \"mhlo.broadcast_in_dim\"(%2958) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2961 = \"mhlo.broadcast_in_dim\"(%2959) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2962 = \"mhlo.concatenate\"(%2960, %2961) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2963 = \"mhlo.broadcast_in_dim\"(%2957) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2964 = \"mhlo.broadcast_in_dim\"(%2962) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2965 = \"mhlo.concatenate\"(%2963, %2964) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2966 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2967 = mhlo.constant dense<0.916666686> : tensor<f32>\n",
      "    %2968 = mhlo.convert %2966 : tensor<f32>\n",
      "    %2969 = mhlo.convert %2967 : tensor<f32>\n",
      "    %2970 = mhlo.power %2968, %2969 : tensor<f32>\n",
      "    %2971 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2972 = mhlo.divide %2971, %2970 : tensor<f32>\n",
      "    %2973 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2974 = mhlo.multiply %2973, %2972 : tensor<f32>\n",
      "    %2975 = mhlo.cosine %2974 : tensor<f32>\n",
      "    %2976 = mhlo.sine %2974 : tensor<f32>\n",
      "    %2977 = mhlo.negate %2976 : tensor<f32>\n",
      "    %2978 = mhlo.convert %2975 : tensor<f32>\n",
      "    %2979 = mhlo.convert %2977 : tensor<f32>\n",
      "    %2980 = \"mhlo.broadcast_in_dim\"(%2978) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2981 = \"mhlo.broadcast_in_dim\"(%2979) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2982 = \"mhlo.concatenate\"(%2980, %2981) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2983 = mhlo.convert %2976 : tensor<f32>\n",
      "    %2984 = mhlo.convert %2975 : tensor<f32>\n",
      "    %2985 = \"mhlo.broadcast_in_dim\"(%2983) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2986 = \"mhlo.broadcast_in_dim\"(%2984) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %2987 = \"mhlo.concatenate\"(%2985, %2986) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %2988 = \"mhlo.broadcast_in_dim\"(%2982) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2989 = \"mhlo.broadcast_in_dim\"(%2987) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %2990 = \"mhlo.concatenate\"(%2988, %2989) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %2991 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %2992 = mhlo.constant dense<0.958333313> : tensor<f32>\n",
      "    %2993 = mhlo.convert %2991 : tensor<f32>\n",
      "    %2994 = mhlo.convert %2992 : tensor<f32>\n",
      "    %2995 = mhlo.power %2993, %2994 : tensor<f32>\n",
      "    %2996 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %2997 = mhlo.divide %2996, %2995 : tensor<f32>\n",
      "    %2998 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %2999 = mhlo.multiply %2998, %2997 : tensor<f32>\n",
      "    %3000 = mhlo.cosine %2999 : tensor<f32>\n",
      "    %3001 = mhlo.sine %2999 : tensor<f32>\n",
      "    %3002 = mhlo.negate %3001 : tensor<f32>\n",
      "    %3003 = mhlo.convert %3000 : tensor<f32>\n",
      "    %3004 = mhlo.convert %3002 : tensor<f32>\n",
      "    %3005 = \"mhlo.broadcast_in_dim\"(%3003) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3006 = \"mhlo.broadcast_in_dim\"(%3004) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3007 = \"mhlo.concatenate\"(%3005, %3006) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3008 = mhlo.convert %3001 : tensor<f32>\n",
      "    %3009 = mhlo.convert %3000 : tensor<f32>\n",
      "    %3010 = \"mhlo.broadcast_in_dim\"(%3008) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3011 = \"mhlo.broadcast_in_dim\"(%3009) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3012 = \"mhlo.concatenate\"(%3010, %3011) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3013 = \"mhlo.broadcast_in_dim\"(%3007) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3014 = \"mhlo.broadcast_in_dim\"(%3012) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3015 = \"mhlo.concatenate\"(%3013, %3014) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3016 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3017 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3018 = mhlo.convert %3016 : tensor<f32>\n",
      "    %3019 = mhlo.convert %3017 : tensor<f32>\n",
      "    %3020 = mhlo.power %3018, %3019 : tensor<f32>\n",
      "    %3021 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3022 = mhlo.divide %3021, %3020 : tensor<f32>\n",
      "    %3023 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3024 = mhlo.multiply %3023, %3022 : tensor<f32>\n",
      "    %3025 = mhlo.cosine %3024 : tensor<f32>\n",
      "    %3026 = mhlo.sine %3024 : tensor<f32>\n",
      "    %3027 = mhlo.negate %3026 : tensor<f32>\n",
      "    %3028 = mhlo.convert %3025 : tensor<f32>\n",
      "    %3029 = mhlo.convert %3027 : tensor<f32>\n",
      "    %3030 = \"mhlo.broadcast_in_dim\"(%3028) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3031 = \"mhlo.broadcast_in_dim\"(%3029) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3032 = \"mhlo.concatenate\"(%3030, %3031) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3033 = mhlo.convert %3026 : tensor<f32>\n",
      "    %3034 = mhlo.convert %3025 : tensor<f32>\n",
      "    %3035 = \"mhlo.broadcast_in_dim\"(%3033) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3036 = \"mhlo.broadcast_in_dim\"(%3034) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3037 = \"mhlo.concatenate\"(%3035, %3036) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3038 = \"mhlo.broadcast_in_dim\"(%3032) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3039 = \"mhlo.broadcast_in_dim\"(%3037) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3040 = \"mhlo.concatenate\"(%3038, %3039) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3041 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3042 = mhlo.constant dense<0.0416666679> : tensor<f32>\n",
      "    %3043 = mhlo.convert %3041 : tensor<f32>\n",
      "    %3044 = mhlo.convert %3042 : tensor<f32>\n",
      "    %3045 = mhlo.power %3043, %3044 : tensor<f32>\n",
      "    %3046 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3047 = mhlo.divide %3046, %3045 : tensor<f32>\n",
      "    %3048 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3049 = mhlo.multiply %3048, %3047 : tensor<f32>\n",
      "    %3050 = mhlo.cosine %3049 : tensor<f32>\n",
      "    %3051 = mhlo.sine %3049 : tensor<f32>\n",
      "    %3052 = mhlo.negate %3051 : tensor<f32>\n",
      "    %3053 = mhlo.convert %3050 : tensor<f32>\n",
      "    %3054 = mhlo.convert %3052 : tensor<f32>\n",
      "    %3055 = \"mhlo.broadcast_in_dim\"(%3053) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3056 = \"mhlo.broadcast_in_dim\"(%3054) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3057 = \"mhlo.concatenate\"(%3055, %3056) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3058 = mhlo.convert %3051 : tensor<f32>\n",
      "    %3059 = mhlo.convert %3050 : tensor<f32>\n",
      "    %3060 = \"mhlo.broadcast_in_dim\"(%3058) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3061 = \"mhlo.broadcast_in_dim\"(%3059) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3062 = \"mhlo.concatenate\"(%3060, %3061) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3063 = \"mhlo.broadcast_in_dim\"(%3057) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3064 = \"mhlo.broadcast_in_dim\"(%3062) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3065 = \"mhlo.concatenate\"(%3063, %3064) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3066 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3067 = mhlo.constant dense<0.0833333358> : tensor<f32>\n",
      "    %3068 = mhlo.convert %3066 : tensor<f32>\n",
      "    %3069 = mhlo.convert %3067 : tensor<f32>\n",
      "    %3070 = mhlo.power %3068, %3069 : tensor<f32>\n",
      "    %3071 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3072 = mhlo.divide %3071, %3070 : tensor<f32>\n",
      "    %3073 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3074 = mhlo.multiply %3073, %3072 : tensor<f32>\n",
      "    %3075 = mhlo.cosine %3074 : tensor<f32>\n",
      "    %3076 = mhlo.sine %3074 : tensor<f32>\n",
      "    %3077 = mhlo.negate %3076 : tensor<f32>\n",
      "    %3078 = mhlo.convert %3075 : tensor<f32>\n",
      "    %3079 = mhlo.convert %3077 : tensor<f32>\n",
      "    %3080 = \"mhlo.broadcast_in_dim\"(%3078) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3081 = \"mhlo.broadcast_in_dim\"(%3079) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3082 = \"mhlo.concatenate\"(%3080, %3081) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3083 = mhlo.convert %3076 : tensor<f32>\n",
      "    %3084 = mhlo.convert %3075 : tensor<f32>\n",
      "    %3085 = \"mhlo.broadcast_in_dim\"(%3083) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3086 = \"mhlo.broadcast_in_dim\"(%3084) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3087 = \"mhlo.concatenate\"(%3085, %3086) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3088 = \"mhlo.broadcast_in_dim\"(%3082) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3089 = \"mhlo.broadcast_in_dim\"(%3087) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3090 = \"mhlo.concatenate\"(%3088, %3089) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3091 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3092 = mhlo.constant dense<1.250000e-01> : tensor<f32>\n",
      "    %3093 = mhlo.convert %3091 : tensor<f32>\n",
      "    %3094 = mhlo.convert %3092 : tensor<f32>\n",
      "    %3095 = mhlo.power %3093, %3094 : tensor<f32>\n",
      "    %3096 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3097 = mhlo.divide %3096, %3095 : tensor<f32>\n",
      "    %3098 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3099 = mhlo.multiply %3098, %3097 : tensor<f32>\n",
      "    %3100 = mhlo.cosine %3099 : tensor<f32>\n",
      "    %3101 = mhlo.sine %3099 : tensor<f32>\n",
      "    %3102 = mhlo.negate %3101 : tensor<f32>\n",
      "    %3103 = mhlo.convert %3100 : tensor<f32>\n",
      "    %3104 = mhlo.convert %3102 : tensor<f32>\n",
      "    %3105 = \"mhlo.broadcast_in_dim\"(%3103) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3106 = \"mhlo.broadcast_in_dim\"(%3104) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3107 = \"mhlo.concatenate\"(%3105, %3106) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3108 = mhlo.convert %3101 : tensor<f32>\n",
      "    %3109 = mhlo.convert %3100 : tensor<f32>\n",
      "    %3110 = \"mhlo.broadcast_in_dim\"(%3108) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3111 = \"mhlo.broadcast_in_dim\"(%3109) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3112 = \"mhlo.concatenate\"(%3110, %3111) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3113 = \"mhlo.broadcast_in_dim\"(%3107) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3114 = \"mhlo.broadcast_in_dim\"(%3112) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3115 = \"mhlo.concatenate\"(%3113, %3114) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3116 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3117 = mhlo.constant dense<0.166666672> : tensor<f32>\n",
      "    %3118 = mhlo.convert %3116 : tensor<f32>\n",
      "    %3119 = mhlo.convert %3117 : tensor<f32>\n",
      "    %3120 = mhlo.power %3118, %3119 : tensor<f32>\n",
      "    %3121 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3122 = mhlo.divide %3121, %3120 : tensor<f32>\n",
      "    %3123 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3124 = mhlo.multiply %3123, %3122 : tensor<f32>\n",
      "    %3125 = mhlo.cosine %3124 : tensor<f32>\n",
      "    %3126 = mhlo.sine %3124 : tensor<f32>\n",
      "    %3127 = mhlo.negate %3126 : tensor<f32>\n",
      "    %3128 = mhlo.convert %3125 : tensor<f32>\n",
      "    %3129 = mhlo.convert %3127 : tensor<f32>\n",
      "    %3130 = \"mhlo.broadcast_in_dim\"(%3128) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3131 = \"mhlo.broadcast_in_dim\"(%3129) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3132 = \"mhlo.concatenate\"(%3130, %3131) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3133 = mhlo.convert %3126 : tensor<f32>\n",
      "    %3134 = mhlo.convert %3125 : tensor<f32>\n",
      "    %3135 = \"mhlo.broadcast_in_dim\"(%3133) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3136 = \"mhlo.broadcast_in_dim\"(%3134) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3137 = \"mhlo.concatenate\"(%3135, %3136) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3138 = \"mhlo.broadcast_in_dim\"(%3132) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3139 = \"mhlo.broadcast_in_dim\"(%3137) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3140 = \"mhlo.concatenate\"(%3138, %3139) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3141 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3142 = mhlo.constant dense<0.208333328> : tensor<f32>\n",
      "    %3143 = mhlo.convert %3141 : tensor<f32>\n",
      "    %3144 = mhlo.convert %3142 : tensor<f32>\n",
      "    %3145 = mhlo.power %3143, %3144 : tensor<f32>\n",
      "    %3146 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3147 = mhlo.divide %3146, %3145 : tensor<f32>\n",
      "    %3148 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3149 = mhlo.multiply %3148, %3147 : tensor<f32>\n",
      "    %3150 = mhlo.cosine %3149 : tensor<f32>\n",
      "    %3151 = mhlo.sine %3149 : tensor<f32>\n",
      "    %3152 = mhlo.negate %3151 : tensor<f32>\n",
      "    %3153 = mhlo.convert %3150 : tensor<f32>\n",
      "    %3154 = mhlo.convert %3152 : tensor<f32>\n",
      "    %3155 = \"mhlo.broadcast_in_dim\"(%3153) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3156 = \"mhlo.broadcast_in_dim\"(%3154) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3157 = \"mhlo.concatenate\"(%3155, %3156) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3158 = mhlo.convert %3151 : tensor<f32>\n",
      "    %3159 = mhlo.convert %3150 : tensor<f32>\n",
      "    %3160 = \"mhlo.broadcast_in_dim\"(%3158) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3161 = \"mhlo.broadcast_in_dim\"(%3159) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3162 = \"mhlo.concatenate\"(%3160, %3161) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3163 = \"mhlo.broadcast_in_dim\"(%3157) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3164 = \"mhlo.broadcast_in_dim\"(%3162) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3165 = \"mhlo.concatenate\"(%3163, %3164) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3166 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3167 = mhlo.constant dense<2.500000e-01> : tensor<f32>\n",
      "    %3168 = mhlo.convert %3166 : tensor<f32>\n",
      "    %3169 = mhlo.convert %3167 : tensor<f32>\n",
      "    %3170 = mhlo.power %3168, %3169 : tensor<f32>\n",
      "    %3171 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3172 = mhlo.divide %3171, %3170 : tensor<f32>\n",
      "    %3173 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3174 = mhlo.multiply %3173, %3172 : tensor<f32>\n",
      "    %3175 = mhlo.cosine %3174 : tensor<f32>\n",
      "    %3176 = mhlo.sine %3174 : tensor<f32>\n",
      "    %3177 = mhlo.negate %3176 : tensor<f32>\n",
      "    %3178 = mhlo.convert %3175 : tensor<f32>\n",
      "    %3179 = mhlo.convert %3177 : tensor<f32>\n",
      "    %3180 = \"mhlo.broadcast_in_dim\"(%3178) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3181 = \"mhlo.broadcast_in_dim\"(%3179) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3182 = \"mhlo.concatenate\"(%3180, %3181) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3183 = mhlo.convert %3176 : tensor<f32>\n",
      "    %3184 = mhlo.convert %3175 : tensor<f32>\n",
      "    %3185 = \"mhlo.broadcast_in_dim\"(%3183) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3186 = \"mhlo.broadcast_in_dim\"(%3184) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3187 = \"mhlo.concatenate\"(%3185, %3186) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3188 = \"mhlo.broadcast_in_dim\"(%3182) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3189 = \"mhlo.broadcast_in_dim\"(%3187) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3190 = \"mhlo.concatenate\"(%3188, %3189) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3191 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3192 = mhlo.constant dense<0.291666657> : tensor<f32>\n",
      "    %3193 = mhlo.convert %3191 : tensor<f32>\n",
      "    %3194 = mhlo.convert %3192 : tensor<f32>\n",
      "    %3195 = mhlo.power %3193, %3194 : tensor<f32>\n",
      "    %3196 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3197 = mhlo.divide %3196, %3195 : tensor<f32>\n",
      "    %3198 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3199 = mhlo.multiply %3198, %3197 : tensor<f32>\n",
      "    %3200 = mhlo.cosine %3199 : tensor<f32>\n",
      "    %3201 = mhlo.sine %3199 : tensor<f32>\n",
      "    %3202 = mhlo.negate %3201 : tensor<f32>\n",
      "    %3203 = mhlo.convert %3200 : tensor<f32>\n",
      "    %3204 = mhlo.convert %3202 : tensor<f32>\n",
      "    %3205 = \"mhlo.broadcast_in_dim\"(%3203) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3206 = \"mhlo.broadcast_in_dim\"(%3204) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3207 = \"mhlo.concatenate\"(%3205, %3206) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3208 = mhlo.convert %3201 : tensor<f32>\n",
      "    %3209 = mhlo.convert %3200 : tensor<f32>\n",
      "    %3210 = \"mhlo.broadcast_in_dim\"(%3208) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3211 = \"mhlo.broadcast_in_dim\"(%3209) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3212 = \"mhlo.concatenate\"(%3210, %3211) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3213 = \"mhlo.broadcast_in_dim\"(%3207) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3214 = \"mhlo.broadcast_in_dim\"(%3212) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3215 = \"mhlo.concatenate\"(%3213, %3214) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3216 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3217 = mhlo.constant dense<0.333333343> : tensor<f32>\n",
      "    %3218 = mhlo.convert %3216 : tensor<f32>\n",
      "    %3219 = mhlo.convert %3217 : tensor<f32>\n",
      "    %3220 = mhlo.power %3218, %3219 : tensor<f32>\n",
      "    %3221 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3222 = mhlo.divide %3221, %3220 : tensor<f32>\n",
      "    %3223 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3224 = mhlo.multiply %3223, %3222 : tensor<f32>\n",
      "    %3225 = mhlo.cosine %3224 : tensor<f32>\n",
      "    %3226 = mhlo.sine %3224 : tensor<f32>\n",
      "    %3227 = mhlo.negate %3226 : tensor<f32>\n",
      "    %3228 = mhlo.convert %3225 : tensor<f32>\n",
      "    %3229 = mhlo.convert %3227 : tensor<f32>\n",
      "    %3230 = \"mhlo.broadcast_in_dim\"(%3228) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3231 = \"mhlo.broadcast_in_dim\"(%3229) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3232 = \"mhlo.concatenate\"(%3230, %3231) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3233 = mhlo.convert %3226 : tensor<f32>\n",
      "    %3234 = mhlo.convert %3225 : tensor<f32>\n",
      "    %3235 = \"mhlo.broadcast_in_dim\"(%3233) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3236 = \"mhlo.broadcast_in_dim\"(%3234) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3237 = \"mhlo.concatenate\"(%3235, %3236) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3238 = \"mhlo.broadcast_in_dim\"(%3232) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3239 = \"mhlo.broadcast_in_dim\"(%3237) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3240 = \"mhlo.concatenate\"(%3238, %3239) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3241 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3242 = mhlo.constant dense<3.750000e-01> : tensor<f32>\n",
      "    %3243 = mhlo.convert %3241 : tensor<f32>\n",
      "    %3244 = mhlo.convert %3242 : tensor<f32>\n",
      "    %3245 = mhlo.power %3243, %3244 : tensor<f32>\n",
      "    %3246 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3247 = mhlo.divide %3246, %3245 : tensor<f32>\n",
      "    %3248 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3249 = mhlo.multiply %3248, %3247 : tensor<f32>\n",
      "    %3250 = mhlo.cosine %3249 : tensor<f32>\n",
      "    %3251 = mhlo.sine %3249 : tensor<f32>\n",
      "    %3252 = mhlo.negate %3251 : tensor<f32>\n",
      "    %3253 = mhlo.convert %3250 : tensor<f32>\n",
      "    %3254 = mhlo.convert %3252 : tensor<f32>\n",
      "    %3255 = \"mhlo.broadcast_in_dim\"(%3253) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3256 = \"mhlo.broadcast_in_dim\"(%3254) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3257 = \"mhlo.concatenate\"(%3255, %3256) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3258 = mhlo.convert %3251 : tensor<f32>\n",
      "    %3259 = mhlo.convert %3250 : tensor<f32>\n",
      "    %3260 = \"mhlo.broadcast_in_dim\"(%3258) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3261 = \"mhlo.broadcast_in_dim\"(%3259) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3262 = \"mhlo.concatenate\"(%3260, %3261) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3263 = \"mhlo.broadcast_in_dim\"(%3257) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3264 = \"mhlo.broadcast_in_dim\"(%3262) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3265 = \"mhlo.concatenate\"(%3263, %3264) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3266 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3267 = mhlo.constant dense<0.416666657> : tensor<f32>\n",
      "    %3268 = mhlo.convert %3266 : tensor<f32>\n",
      "    %3269 = mhlo.convert %3267 : tensor<f32>\n",
      "    %3270 = mhlo.power %3268, %3269 : tensor<f32>\n",
      "    %3271 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3272 = mhlo.divide %3271, %3270 : tensor<f32>\n",
      "    %3273 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3274 = mhlo.multiply %3273, %3272 : tensor<f32>\n",
      "    %3275 = mhlo.cosine %3274 : tensor<f32>\n",
      "    %3276 = mhlo.sine %3274 : tensor<f32>\n",
      "    %3277 = mhlo.negate %3276 : tensor<f32>\n",
      "    %3278 = mhlo.convert %3275 : tensor<f32>\n",
      "    %3279 = mhlo.convert %3277 : tensor<f32>\n",
      "    %3280 = \"mhlo.broadcast_in_dim\"(%3278) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3281 = \"mhlo.broadcast_in_dim\"(%3279) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3282 = \"mhlo.concatenate\"(%3280, %3281) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3283 = mhlo.convert %3276 : tensor<f32>\n",
      "    %3284 = mhlo.convert %3275 : tensor<f32>\n",
      "    %3285 = \"mhlo.broadcast_in_dim\"(%3283) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3286 = \"mhlo.broadcast_in_dim\"(%3284) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3287 = \"mhlo.concatenate\"(%3285, %3286) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3288 = \"mhlo.broadcast_in_dim\"(%3282) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3289 = \"mhlo.broadcast_in_dim\"(%3287) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3290 = \"mhlo.concatenate\"(%3288, %3289) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3291 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3292 = mhlo.constant dense<0.458333343> : tensor<f32>\n",
      "    %3293 = mhlo.convert %3291 : tensor<f32>\n",
      "    %3294 = mhlo.convert %3292 : tensor<f32>\n",
      "    %3295 = mhlo.power %3293, %3294 : tensor<f32>\n",
      "    %3296 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3297 = mhlo.divide %3296, %3295 : tensor<f32>\n",
      "    %3298 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3299 = mhlo.multiply %3298, %3297 : tensor<f32>\n",
      "    %3300 = mhlo.cosine %3299 : tensor<f32>\n",
      "    %3301 = mhlo.sine %3299 : tensor<f32>\n",
      "    %3302 = mhlo.negate %3301 : tensor<f32>\n",
      "    %3303 = mhlo.convert %3300 : tensor<f32>\n",
      "    %3304 = mhlo.convert %3302 : tensor<f32>\n",
      "    %3305 = \"mhlo.broadcast_in_dim\"(%3303) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3306 = \"mhlo.broadcast_in_dim\"(%3304) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3307 = \"mhlo.concatenate\"(%3305, %3306) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3308 = mhlo.convert %3301 : tensor<f32>\n",
      "    %3309 = mhlo.convert %3300 : tensor<f32>\n",
      "    %3310 = \"mhlo.broadcast_in_dim\"(%3308) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3311 = \"mhlo.broadcast_in_dim\"(%3309) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3312 = \"mhlo.concatenate\"(%3310, %3311) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3313 = \"mhlo.broadcast_in_dim\"(%3307) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3314 = \"mhlo.broadcast_in_dim\"(%3312) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3315 = \"mhlo.concatenate\"(%3313, %3314) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3316 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3317 = mhlo.constant dense<5.000000e-01> : tensor<f32>\n",
      "    %3318 = mhlo.convert %3316 : tensor<f32>\n",
      "    %3319 = mhlo.convert %3317 : tensor<f32>\n",
      "    %3320 = mhlo.power %3318, %3319 : tensor<f32>\n",
      "    %3321 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3322 = mhlo.divide %3321, %3320 : tensor<f32>\n",
      "    %3323 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3324 = mhlo.multiply %3323, %3322 : tensor<f32>\n",
      "    %3325 = mhlo.cosine %3324 : tensor<f32>\n",
      "    %3326 = mhlo.sine %3324 : tensor<f32>\n",
      "    %3327 = mhlo.negate %3326 : tensor<f32>\n",
      "    %3328 = mhlo.convert %3325 : tensor<f32>\n",
      "    %3329 = mhlo.convert %3327 : tensor<f32>\n",
      "    %3330 = \"mhlo.broadcast_in_dim\"(%3328) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3331 = \"mhlo.broadcast_in_dim\"(%3329) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3332 = \"mhlo.concatenate\"(%3330, %3331) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3333 = mhlo.convert %3326 : tensor<f32>\n",
      "    %3334 = mhlo.convert %3325 : tensor<f32>\n",
      "    %3335 = \"mhlo.broadcast_in_dim\"(%3333) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3336 = \"mhlo.broadcast_in_dim\"(%3334) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3337 = \"mhlo.concatenate\"(%3335, %3336) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3338 = \"mhlo.broadcast_in_dim\"(%3332) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3339 = \"mhlo.broadcast_in_dim\"(%3337) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3340 = \"mhlo.concatenate\"(%3338, %3339) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3341 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3342 = mhlo.constant dense<0.541666687> : tensor<f32>\n",
      "    %3343 = mhlo.convert %3341 : tensor<f32>\n",
      "    %3344 = mhlo.convert %3342 : tensor<f32>\n",
      "    %3345 = mhlo.power %3343, %3344 : tensor<f32>\n",
      "    %3346 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3347 = mhlo.divide %3346, %3345 : tensor<f32>\n",
      "    %3348 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3349 = mhlo.multiply %3348, %3347 : tensor<f32>\n",
      "    %3350 = mhlo.cosine %3349 : tensor<f32>\n",
      "    %3351 = mhlo.sine %3349 : tensor<f32>\n",
      "    %3352 = mhlo.negate %3351 : tensor<f32>\n",
      "    %3353 = mhlo.convert %3350 : tensor<f32>\n",
      "    %3354 = mhlo.convert %3352 : tensor<f32>\n",
      "    %3355 = \"mhlo.broadcast_in_dim\"(%3353) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3356 = \"mhlo.broadcast_in_dim\"(%3354) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3357 = \"mhlo.concatenate\"(%3355, %3356) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3358 = mhlo.convert %3351 : tensor<f32>\n",
      "    %3359 = mhlo.convert %3350 : tensor<f32>\n",
      "    %3360 = \"mhlo.broadcast_in_dim\"(%3358) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3361 = \"mhlo.broadcast_in_dim\"(%3359) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3362 = \"mhlo.concatenate\"(%3360, %3361) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3363 = \"mhlo.broadcast_in_dim\"(%3357) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3364 = \"mhlo.broadcast_in_dim\"(%3362) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3365 = \"mhlo.concatenate\"(%3363, %3364) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3366 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3367 = mhlo.constant dense<0.583333313> : tensor<f32>\n",
      "    %3368 = mhlo.convert %3366 : tensor<f32>\n",
      "    %3369 = mhlo.convert %3367 : tensor<f32>\n",
      "    %3370 = mhlo.power %3368, %3369 : tensor<f32>\n",
      "    %3371 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3372 = mhlo.divide %3371, %3370 : tensor<f32>\n",
      "    %3373 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3374 = mhlo.multiply %3373, %3372 : tensor<f32>\n",
      "    %3375 = mhlo.cosine %3374 : tensor<f32>\n",
      "    %3376 = mhlo.sine %3374 : tensor<f32>\n",
      "    %3377 = mhlo.negate %3376 : tensor<f32>\n",
      "    %3378 = mhlo.convert %3375 : tensor<f32>\n",
      "    %3379 = mhlo.convert %3377 : tensor<f32>\n",
      "    %3380 = \"mhlo.broadcast_in_dim\"(%3378) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3381 = \"mhlo.broadcast_in_dim\"(%3379) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3382 = \"mhlo.concatenate\"(%3380, %3381) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3383 = mhlo.convert %3376 : tensor<f32>\n",
      "    %3384 = mhlo.convert %3375 : tensor<f32>\n",
      "    %3385 = \"mhlo.broadcast_in_dim\"(%3383) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3386 = \"mhlo.broadcast_in_dim\"(%3384) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3387 = \"mhlo.concatenate\"(%3385, %3386) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3388 = \"mhlo.broadcast_in_dim\"(%3382) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3389 = \"mhlo.broadcast_in_dim\"(%3387) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3390 = \"mhlo.concatenate\"(%3388, %3389) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3391 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3392 = mhlo.constant dense<6.250000e-01> : tensor<f32>\n",
      "    %3393 = mhlo.convert %3391 : tensor<f32>\n",
      "    %3394 = mhlo.convert %3392 : tensor<f32>\n",
      "    %3395 = mhlo.power %3393, %3394 : tensor<f32>\n",
      "    %3396 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3397 = mhlo.divide %3396, %3395 : tensor<f32>\n",
      "    %3398 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3399 = mhlo.multiply %3398, %3397 : tensor<f32>\n",
      "    %3400 = mhlo.cosine %3399 : tensor<f32>\n",
      "    %3401 = mhlo.sine %3399 : tensor<f32>\n",
      "    %3402 = mhlo.negate %3401 : tensor<f32>\n",
      "    %3403 = mhlo.convert %3400 : tensor<f32>\n",
      "    %3404 = mhlo.convert %3402 : tensor<f32>\n",
      "    %3405 = \"mhlo.broadcast_in_dim\"(%3403) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3406 = \"mhlo.broadcast_in_dim\"(%3404) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3407 = \"mhlo.concatenate\"(%3405, %3406) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3408 = mhlo.convert %3401 : tensor<f32>\n",
      "    %3409 = mhlo.convert %3400 : tensor<f32>\n",
      "    %3410 = \"mhlo.broadcast_in_dim\"(%3408) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3411 = \"mhlo.broadcast_in_dim\"(%3409) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3412 = \"mhlo.concatenate\"(%3410, %3411) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3413 = \"mhlo.broadcast_in_dim\"(%3407) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3414 = \"mhlo.broadcast_in_dim\"(%3412) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3415 = \"mhlo.concatenate\"(%3413, %3414) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3416 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3417 = mhlo.constant dense<0.666666686> : tensor<f32>\n",
      "    %3418 = mhlo.convert %3416 : tensor<f32>\n",
      "    %3419 = mhlo.convert %3417 : tensor<f32>\n",
      "    %3420 = mhlo.power %3418, %3419 : tensor<f32>\n",
      "    %3421 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3422 = mhlo.divide %3421, %3420 : tensor<f32>\n",
      "    %3423 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3424 = mhlo.multiply %3423, %3422 : tensor<f32>\n",
      "    %3425 = mhlo.cosine %3424 : tensor<f32>\n",
      "    %3426 = mhlo.sine %3424 : tensor<f32>\n",
      "    %3427 = mhlo.negate %3426 : tensor<f32>\n",
      "    %3428 = mhlo.convert %3425 : tensor<f32>\n",
      "    %3429 = mhlo.convert %3427 : tensor<f32>\n",
      "    %3430 = \"mhlo.broadcast_in_dim\"(%3428) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3431 = \"mhlo.broadcast_in_dim\"(%3429) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3432 = \"mhlo.concatenate\"(%3430, %3431) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3433 = mhlo.convert %3426 : tensor<f32>\n",
      "    %3434 = mhlo.convert %3425 : tensor<f32>\n",
      "    %3435 = \"mhlo.broadcast_in_dim\"(%3433) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3436 = \"mhlo.broadcast_in_dim\"(%3434) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3437 = \"mhlo.concatenate\"(%3435, %3436) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3438 = \"mhlo.broadcast_in_dim\"(%3432) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3439 = \"mhlo.broadcast_in_dim\"(%3437) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3440 = \"mhlo.concatenate\"(%3438, %3439) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3441 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3442 = mhlo.constant dense<0.708333313> : tensor<f32>\n",
      "    %3443 = mhlo.convert %3441 : tensor<f32>\n",
      "    %3444 = mhlo.convert %3442 : tensor<f32>\n",
      "    %3445 = mhlo.power %3443, %3444 : tensor<f32>\n",
      "    %3446 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3447 = mhlo.divide %3446, %3445 : tensor<f32>\n",
      "    %3448 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3449 = mhlo.multiply %3448, %3447 : tensor<f32>\n",
      "    %3450 = mhlo.cosine %3449 : tensor<f32>\n",
      "    %3451 = mhlo.sine %3449 : tensor<f32>\n",
      "    %3452 = mhlo.negate %3451 : tensor<f32>\n",
      "    %3453 = mhlo.convert %3450 : tensor<f32>\n",
      "    %3454 = mhlo.convert %3452 : tensor<f32>\n",
      "    %3455 = \"mhlo.broadcast_in_dim\"(%3453) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3456 = \"mhlo.broadcast_in_dim\"(%3454) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3457 = \"mhlo.concatenate\"(%3455, %3456) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3458 = mhlo.convert %3451 : tensor<f32>\n",
      "    %3459 = mhlo.convert %3450 : tensor<f32>\n",
      "    %3460 = \"mhlo.broadcast_in_dim\"(%3458) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3461 = \"mhlo.broadcast_in_dim\"(%3459) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3462 = \"mhlo.concatenate\"(%3460, %3461) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3463 = \"mhlo.broadcast_in_dim\"(%3457) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3464 = \"mhlo.broadcast_in_dim\"(%3462) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3465 = \"mhlo.concatenate\"(%3463, %3464) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3466 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3467 = mhlo.constant dense<7.500000e-01> : tensor<f32>\n",
      "    %3468 = mhlo.convert %3466 : tensor<f32>\n",
      "    %3469 = mhlo.convert %3467 : tensor<f32>\n",
      "    %3470 = mhlo.power %3468, %3469 : tensor<f32>\n",
      "    %3471 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3472 = mhlo.divide %3471, %3470 : tensor<f32>\n",
      "    %3473 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3474 = mhlo.multiply %3473, %3472 : tensor<f32>\n",
      "    %3475 = mhlo.cosine %3474 : tensor<f32>\n",
      "    %3476 = mhlo.sine %3474 : tensor<f32>\n",
      "    %3477 = mhlo.negate %3476 : tensor<f32>\n",
      "    %3478 = mhlo.convert %3475 : tensor<f32>\n",
      "    %3479 = mhlo.convert %3477 : tensor<f32>\n",
      "    %3480 = \"mhlo.broadcast_in_dim\"(%3478) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3481 = \"mhlo.broadcast_in_dim\"(%3479) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3482 = \"mhlo.concatenate\"(%3480, %3481) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3483 = mhlo.convert %3476 : tensor<f32>\n",
      "    %3484 = mhlo.convert %3475 : tensor<f32>\n",
      "    %3485 = \"mhlo.broadcast_in_dim\"(%3483) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3486 = \"mhlo.broadcast_in_dim\"(%3484) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3487 = \"mhlo.concatenate\"(%3485, %3486) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3488 = \"mhlo.broadcast_in_dim\"(%3482) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3489 = \"mhlo.broadcast_in_dim\"(%3487) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3490 = \"mhlo.concatenate\"(%3488, %3489) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3491 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3492 = mhlo.constant dense<0.791666686> : tensor<f32>\n",
      "    %3493 = mhlo.convert %3491 : tensor<f32>\n",
      "    %3494 = mhlo.convert %3492 : tensor<f32>\n",
      "    %3495 = mhlo.power %3493, %3494 : tensor<f32>\n",
      "    %3496 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3497 = mhlo.divide %3496, %3495 : tensor<f32>\n",
      "    %3498 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3499 = mhlo.multiply %3498, %3497 : tensor<f32>\n",
      "    %3500 = mhlo.cosine %3499 : tensor<f32>\n",
      "    %3501 = mhlo.sine %3499 : tensor<f32>\n",
      "    %3502 = mhlo.negate %3501 : tensor<f32>\n",
      "    %3503 = mhlo.convert %3500 : tensor<f32>\n",
      "    %3504 = mhlo.convert %3502 : tensor<f32>\n",
      "    %3505 = \"mhlo.broadcast_in_dim\"(%3503) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3506 = \"mhlo.broadcast_in_dim\"(%3504) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3507 = \"mhlo.concatenate\"(%3505, %3506) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3508 = mhlo.convert %3501 : tensor<f32>\n",
      "    %3509 = mhlo.convert %3500 : tensor<f32>\n",
      "    %3510 = \"mhlo.broadcast_in_dim\"(%3508) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3511 = \"mhlo.broadcast_in_dim\"(%3509) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3512 = \"mhlo.concatenate\"(%3510, %3511) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3513 = \"mhlo.broadcast_in_dim\"(%3507) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3514 = \"mhlo.broadcast_in_dim\"(%3512) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3515 = \"mhlo.concatenate\"(%3513, %3514) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3516 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3517 = mhlo.constant dense<0.833333313> : tensor<f32>\n",
      "    %3518 = mhlo.convert %3516 : tensor<f32>\n",
      "    %3519 = mhlo.convert %3517 : tensor<f32>\n",
      "    %3520 = mhlo.power %3518, %3519 : tensor<f32>\n",
      "    %3521 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3522 = mhlo.divide %3521, %3520 : tensor<f32>\n",
      "    %3523 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3524 = mhlo.multiply %3523, %3522 : tensor<f32>\n",
      "    %3525 = mhlo.cosine %3524 : tensor<f32>\n",
      "    %3526 = mhlo.sine %3524 : tensor<f32>\n",
      "    %3527 = mhlo.negate %3526 : tensor<f32>\n",
      "    %3528 = mhlo.convert %3525 : tensor<f32>\n",
      "    %3529 = mhlo.convert %3527 : tensor<f32>\n",
      "    %3530 = \"mhlo.broadcast_in_dim\"(%3528) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3531 = \"mhlo.broadcast_in_dim\"(%3529) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3532 = \"mhlo.concatenate\"(%3530, %3531) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3533 = mhlo.convert %3526 : tensor<f32>\n",
      "    %3534 = mhlo.convert %3525 : tensor<f32>\n",
      "    %3535 = \"mhlo.broadcast_in_dim\"(%3533) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3536 = \"mhlo.broadcast_in_dim\"(%3534) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3537 = \"mhlo.concatenate\"(%3535, %3536) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3538 = \"mhlo.broadcast_in_dim\"(%3532) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3539 = \"mhlo.broadcast_in_dim\"(%3537) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3540 = \"mhlo.concatenate\"(%3538, %3539) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3541 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3542 = mhlo.constant dense<8.750000e-01> : tensor<f32>\n",
      "    %3543 = mhlo.convert %3541 : tensor<f32>\n",
      "    %3544 = mhlo.convert %3542 : tensor<f32>\n",
      "    %3545 = mhlo.power %3543, %3544 : tensor<f32>\n",
      "    %3546 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3547 = mhlo.divide %3546, %3545 : tensor<f32>\n",
      "    %3548 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3549 = mhlo.multiply %3548, %3547 : tensor<f32>\n",
      "    %3550 = mhlo.cosine %3549 : tensor<f32>\n",
      "    %3551 = mhlo.sine %3549 : tensor<f32>\n",
      "    %3552 = mhlo.negate %3551 : tensor<f32>\n",
      "    %3553 = mhlo.convert %3550 : tensor<f32>\n",
      "    %3554 = mhlo.convert %3552 : tensor<f32>\n",
      "    %3555 = \"mhlo.broadcast_in_dim\"(%3553) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3556 = \"mhlo.broadcast_in_dim\"(%3554) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3557 = \"mhlo.concatenate\"(%3555, %3556) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3558 = mhlo.convert %3551 : tensor<f32>\n",
      "    %3559 = mhlo.convert %3550 : tensor<f32>\n",
      "    %3560 = \"mhlo.broadcast_in_dim\"(%3558) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3561 = \"mhlo.broadcast_in_dim\"(%3559) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3562 = \"mhlo.concatenate\"(%3560, %3561) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3563 = \"mhlo.broadcast_in_dim\"(%3557) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3564 = \"mhlo.broadcast_in_dim\"(%3562) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3565 = \"mhlo.concatenate\"(%3563, %3564) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3566 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3567 = mhlo.constant dense<0.916666686> : tensor<f32>\n",
      "    %3568 = mhlo.convert %3566 : tensor<f32>\n",
      "    %3569 = mhlo.convert %3567 : tensor<f32>\n",
      "    %3570 = mhlo.power %3568, %3569 : tensor<f32>\n",
      "    %3571 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3572 = mhlo.divide %3571, %3570 : tensor<f32>\n",
      "    %3573 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3574 = mhlo.multiply %3573, %3572 : tensor<f32>\n",
      "    %3575 = mhlo.cosine %3574 : tensor<f32>\n",
      "    %3576 = mhlo.sine %3574 : tensor<f32>\n",
      "    %3577 = mhlo.negate %3576 : tensor<f32>\n",
      "    %3578 = mhlo.convert %3575 : tensor<f32>\n",
      "    %3579 = mhlo.convert %3577 : tensor<f32>\n",
      "    %3580 = \"mhlo.broadcast_in_dim\"(%3578) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3581 = \"mhlo.broadcast_in_dim\"(%3579) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3582 = \"mhlo.concatenate\"(%3580, %3581) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3583 = mhlo.convert %3576 : tensor<f32>\n",
      "    %3584 = mhlo.convert %3575 : tensor<f32>\n",
      "    %3585 = \"mhlo.broadcast_in_dim\"(%3583) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3586 = \"mhlo.broadcast_in_dim\"(%3584) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3587 = \"mhlo.concatenate\"(%3585, %3586) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3588 = \"mhlo.broadcast_in_dim\"(%3582) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3589 = \"mhlo.broadcast_in_dim\"(%3587) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3590 = \"mhlo.concatenate\"(%3588, %3589) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3591 = mhlo.constant dense<1.000000e+04> : tensor<f32>\n",
      "    %3592 = mhlo.constant dense<0.958333313> : tensor<f32>\n",
      "    %3593 = mhlo.convert %3591 : tensor<f32>\n",
      "    %3594 = mhlo.convert %3592 : tensor<f32>\n",
      "    %3595 = mhlo.power %3593, %3594 : tensor<f32>\n",
      "    %3596 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3597 = mhlo.divide %3596, %3595 : tensor<f32>\n",
      "    %3598 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3599 = mhlo.multiply %3598, %3597 : tensor<f32>\n",
      "    %3600 = mhlo.cosine %3599 : tensor<f32>\n",
      "    %3601 = mhlo.sine %3599 : tensor<f32>\n",
      "    %3602 = mhlo.negate %3601 : tensor<f32>\n",
      "    %3603 = mhlo.convert %3600 : tensor<f32>\n",
      "    %3604 = mhlo.convert %3602 : tensor<f32>\n",
      "    %3605 = \"mhlo.broadcast_in_dim\"(%3603) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3606 = \"mhlo.broadcast_in_dim\"(%3604) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3607 = \"mhlo.concatenate\"(%3605, %3606) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3608 = mhlo.convert %3601 : tensor<f32>\n",
      "    %3609 = mhlo.convert %3600 : tensor<f32>\n",
      "    %3610 = \"mhlo.broadcast_in_dim\"(%3608) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3611 = \"mhlo.broadcast_in_dim\"(%3609) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3612 = \"mhlo.concatenate\"(%3610, %3611) {dimension = 0 : i64} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>\n",
      "    %3613 = \"mhlo.broadcast_in_dim\"(%3607) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3614 = \"mhlo.broadcast_in_dim\"(%3612) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<2xf32>) -> tensor<1x2xf32>\n",
      "    %3615 = \"mhlo.concatenate\"(%3613, %3614) {dimension = 0 : i64} : (tensor<1x2xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>\n",
      "    %3616 = \"mhlo.broadcast_in_dim\"(%40) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3617 = \"mhlo.broadcast_in_dim\"(%65) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3618 = \"mhlo.broadcast_in_dim\"(%90) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3619 = \"mhlo.broadcast_in_dim\"(%115) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3620 = \"mhlo.broadcast_in_dim\"(%140) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3621 = \"mhlo.broadcast_in_dim\"(%165) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3622 = \"mhlo.broadcast_in_dim\"(%190) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3623 = \"mhlo.broadcast_in_dim\"(%215) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3624 = \"mhlo.broadcast_in_dim\"(%240) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3625 = \"mhlo.broadcast_in_dim\"(%265) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3626 = \"mhlo.broadcast_in_dim\"(%290) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3627 = \"mhlo.broadcast_in_dim\"(%315) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3628 = \"mhlo.broadcast_in_dim\"(%340) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3629 = \"mhlo.broadcast_in_dim\"(%365) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3630 = \"mhlo.broadcast_in_dim\"(%390) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3631 = \"mhlo.broadcast_in_dim\"(%415) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3632 = \"mhlo.broadcast_in_dim\"(%440) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3633 = \"mhlo.broadcast_in_dim\"(%465) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3634 = \"mhlo.broadcast_in_dim\"(%490) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3635 = \"mhlo.broadcast_in_dim\"(%515) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3636 = \"mhlo.broadcast_in_dim\"(%540) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3637 = \"mhlo.broadcast_in_dim\"(%565) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3638 = \"mhlo.broadcast_in_dim\"(%590) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3639 = \"mhlo.broadcast_in_dim\"(%615) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3640 = \"mhlo.broadcast_in_dim\"(%640) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3641 = \"mhlo.broadcast_in_dim\"(%665) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3642 = \"mhlo.broadcast_in_dim\"(%690) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3643 = \"mhlo.broadcast_in_dim\"(%715) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3644 = \"mhlo.broadcast_in_dim\"(%740) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3645 = \"mhlo.broadcast_in_dim\"(%765) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3646 = \"mhlo.broadcast_in_dim\"(%790) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3647 = \"mhlo.broadcast_in_dim\"(%815) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3648 = \"mhlo.broadcast_in_dim\"(%840) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3649 = \"mhlo.broadcast_in_dim\"(%865) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3650 = \"mhlo.broadcast_in_dim\"(%890) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3651 = \"mhlo.broadcast_in_dim\"(%915) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3652 = \"mhlo.broadcast_in_dim\"(%940) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3653 = \"mhlo.broadcast_in_dim\"(%965) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3654 = \"mhlo.broadcast_in_dim\"(%990) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3655 = \"mhlo.broadcast_in_dim\"(%1015) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3656 = \"mhlo.broadcast_in_dim\"(%1040) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3657 = \"mhlo.broadcast_in_dim\"(%1065) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3658 = \"mhlo.broadcast_in_dim\"(%1090) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3659 = \"mhlo.broadcast_in_dim\"(%1115) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3660 = \"mhlo.broadcast_in_dim\"(%1140) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3661 = \"mhlo.broadcast_in_dim\"(%1165) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3662 = \"mhlo.broadcast_in_dim\"(%1190) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3663 = \"mhlo.broadcast_in_dim\"(%1215) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3664 = \"mhlo.broadcast_in_dim\"(%1240) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3665 = \"mhlo.broadcast_in_dim\"(%1265) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3666 = \"mhlo.broadcast_in_dim\"(%1290) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3667 = \"mhlo.broadcast_in_dim\"(%1315) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3668 = \"mhlo.broadcast_in_dim\"(%1340) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3669 = \"mhlo.broadcast_in_dim\"(%1365) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3670 = \"mhlo.broadcast_in_dim\"(%1390) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3671 = \"mhlo.broadcast_in_dim\"(%1415) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3672 = \"mhlo.broadcast_in_dim\"(%1440) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3673 = \"mhlo.broadcast_in_dim\"(%1465) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3674 = \"mhlo.broadcast_in_dim\"(%1490) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3675 = \"mhlo.broadcast_in_dim\"(%1515) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3676 = \"mhlo.broadcast_in_dim\"(%1540) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3677 = \"mhlo.broadcast_in_dim\"(%1565) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3678 = \"mhlo.broadcast_in_dim\"(%1590) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3679 = \"mhlo.broadcast_in_dim\"(%1615) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3680 = \"mhlo.broadcast_in_dim\"(%1640) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3681 = \"mhlo.broadcast_in_dim\"(%1665) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3682 = \"mhlo.broadcast_in_dim\"(%1690) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3683 = \"mhlo.broadcast_in_dim\"(%1715) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3684 = \"mhlo.broadcast_in_dim\"(%1740) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3685 = \"mhlo.broadcast_in_dim\"(%1765) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3686 = \"mhlo.broadcast_in_dim\"(%1790) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3687 = \"mhlo.broadcast_in_dim\"(%1815) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3688 = \"mhlo.broadcast_in_dim\"(%1840) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3689 = \"mhlo.broadcast_in_dim\"(%1865) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3690 = \"mhlo.broadcast_in_dim\"(%1890) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3691 = \"mhlo.broadcast_in_dim\"(%1915) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3692 = \"mhlo.broadcast_in_dim\"(%1940) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3693 = \"mhlo.broadcast_in_dim\"(%1965) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3694 = \"mhlo.broadcast_in_dim\"(%1990) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3695 = \"mhlo.broadcast_in_dim\"(%2015) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3696 = \"mhlo.broadcast_in_dim\"(%2040) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3697 = \"mhlo.broadcast_in_dim\"(%2065) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3698 = \"mhlo.broadcast_in_dim\"(%2090) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3699 = \"mhlo.broadcast_in_dim\"(%2115) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3700 = \"mhlo.broadcast_in_dim\"(%2140) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3701 = \"mhlo.broadcast_in_dim\"(%2165) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3702 = \"mhlo.broadcast_in_dim\"(%2190) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3703 = \"mhlo.broadcast_in_dim\"(%2215) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3704 = \"mhlo.broadcast_in_dim\"(%2240) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3705 = \"mhlo.broadcast_in_dim\"(%2265) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3706 = \"mhlo.broadcast_in_dim\"(%2290) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3707 = \"mhlo.broadcast_in_dim\"(%2315) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3708 = \"mhlo.broadcast_in_dim\"(%2340) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3709 = \"mhlo.broadcast_in_dim\"(%2365) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3710 = \"mhlo.broadcast_in_dim\"(%2390) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3711 = \"mhlo.broadcast_in_dim\"(%2415) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3712 = \"mhlo.broadcast_in_dim\"(%2440) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3713 = \"mhlo.broadcast_in_dim\"(%2465) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3714 = \"mhlo.broadcast_in_dim\"(%2490) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3715 = \"mhlo.broadcast_in_dim\"(%2515) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3716 = \"mhlo.broadcast_in_dim\"(%2540) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3717 = \"mhlo.broadcast_in_dim\"(%2565) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3718 = \"mhlo.broadcast_in_dim\"(%2590) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3719 = \"mhlo.broadcast_in_dim\"(%2615) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3720 = \"mhlo.broadcast_in_dim\"(%2640) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3721 = \"mhlo.broadcast_in_dim\"(%2665) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3722 = \"mhlo.broadcast_in_dim\"(%2690) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3723 = \"mhlo.broadcast_in_dim\"(%2715) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3724 = \"mhlo.broadcast_in_dim\"(%2740) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3725 = \"mhlo.broadcast_in_dim\"(%2765) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3726 = \"mhlo.broadcast_in_dim\"(%2790) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3727 = \"mhlo.broadcast_in_dim\"(%2815) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3728 = \"mhlo.broadcast_in_dim\"(%2840) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3729 = \"mhlo.broadcast_in_dim\"(%2865) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3730 = \"mhlo.broadcast_in_dim\"(%2890) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3731 = \"mhlo.broadcast_in_dim\"(%2915) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3732 = \"mhlo.broadcast_in_dim\"(%2940) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3733 = \"mhlo.broadcast_in_dim\"(%2965) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3734 = \"mhlo.broadcast_in_dim\"(%2990) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3735 = \"mhlo.broadcast_in_dim\"(%3015) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3736 = \"mhlo.broadcast_in_dim\"(%3040) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3737 = \"mhlo.broadcast_in_dim\"(%3065) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3738 = \"mhlo.broadcast_in_dim\"(%3090) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3739 = \"mhlo.broadcast_in_dim\"(%3115) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3740 = \"mhlo.broadcast_in_dim\"(%3140) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3741 = \"mhlo.broadcast_in_dim\"(%3165) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3742 = \"mhlo.broadcast_in_dim\"(%3190) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3743 = \"mhlo.broadcast_in_dim\"(%3215) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3744 = \"mhlo.broadcast_in_dim\"(%3240) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3745 = \"mhlo.broadcast_in_dim\"(%3265) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3746 = \"mhlo.broadcast_in_dim\"(%3290) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3747 = \"mhlo.broadcast_in_dim\"(%3315) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3748 = \"mhlo.broadcast_in_dim\"(%3340) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3749 = \"mhlo.broadcast_in_dim\"(%3365) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3750 = \"mhlo.broadcast_in_dim\"(%3390) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3751 = \"mhlo.broadcast_in_dim\"(%3415) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3752 = \"mhlo.broadcast_in_dim\"(%3440) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3753 = \"mhlo.broadcast_in_dim\"(%3465) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3754 = \"mhlo.broadcast_in_dim\"(%3490) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3755 = \"mhlo.broadcast_in_dim\"(%3515) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3756 = \"mhlo.broadcast_in_dim\"(%3540) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3757 = \"mhlo.broadcast_in_dim\"(%3565) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3758 = \"mhlo.broadcast_in_dim\"(%3590) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3759 = \"mhlo.broadcast_in_dim\"(%3615) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3760 = \"mhlo.concatenate\"(%3616, %3617, %3618, %3619, %3620, %3621, %3622, %3623, %3624, %3625, %3626, %3627, %3628, %3629, %3630, %3631) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3761 = \"mhlo.concatenate\"(%3632, %3633, %3634, %3635, %3636, %3637, %3638, %3639, %3640, %3641, %3642, %3643, %3644, %3645, %3646, %3647) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3762 = \"mhlo.concatenate\"(%3648, %3649, %3650, %3651, %3652, %3653, %3654, %3655, %3656, %3657, %3658, %3659, %3660, %3661, %3662, %3663) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3763 = \"mhlo.concatenate\"(%3664, %3665, %3666, %3667, %3668, %3669, %3670, %3671, %3672, %3673, %3674, %3675, %3676, %3677, %3678, %3679) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3764 = \"mhlo.concatenate\"(%3680, %3681, %3682, %3683, %3684, %3685, %3686, %3687, %3688, %3689, %3690, %3691, %3692, %3693, %3694, %3695) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3765 = \"mhlo.concatenate\"(%3696, %3697, %3698, %3699, %3700, %3701, %3702, %3703, %3704, %3705, %3706, %3707, %3708, %3709, %3710, %3711) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3766 = \"mhlo.concatenate\"(%3712, %3713, %3714, %3715, %3716, %3717, %3718, %3719, %3720, %3721, %3722, %3723, %3724, %3725, %3726, %3727) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3767 = \"mhlo.concatenate\"(%3728, %3729, %3730, %3731, %3732, %3733, %3734, %3735, %3736, %3737, %3738, %3739, %3740, %3741, %3742, %3743) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3768 = \"mhlo.concatenate\"(%3744, %3745, %3746, %3747, %3748, %3749, %3750, %3751, %3752, %3753, %3754, %3755, %3756, %3757, %3758, %3759) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3769 = \"mhlo.concatenate\"(%3760, %3761, %3762, %3763, %3764, %3765, %3766, %3767, %3768) {dimension = 0 : i64} : (tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>) -> tensor<144x2x2xf32>\n",
      "    %3770 = \"mhlo.broadcast_in_dim\"(%40) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3771 = \"mhlo.broadcast_in_dim\"(%65) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3772 = \"mhlo.broadcast_in_dim\"(%90) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3773 = \"mhlo.broadcast_in_dim\"(%115) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3774 = \"mhlo.broadcast_in_dim\"(%140) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3775 = \"mhlo.broadcast_in_dim\"(%165) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3776 = \"mhlo.broadcast_in_dim\"(%190) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3777 = \"mhlo.broadcast_in_dim\"(%215) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3778 = \"mhlo.broadcast_in_dim\"(%240) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3779 = \"mhlo.broadcast_in_dim\"(%265) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3780 = \"mhlo.broadcast_in_dim\"(%290) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3781 = \"mhlo.broadcast_in_dim\"(%315) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3782 = \"mhlo.broadcast_in_dim\"(%340) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3783 = \"mhlo.broadcast_in_dim\"(%365) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3784 = \"mhlo.broadcast_in_dim\"(%390) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3785 = \"mhlo.broadcast_in_dim\"(%415) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3786 = \"mhlo.broadcast_in_dim\"(%440) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3787 = \"mhlo.broadcast_in_dim\"(%465) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3788 = \"mhlo.broadcast_in_dim\"(%490) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3789 = \"mhlo.broadcast_in_dim\"(%515) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3790 = \"mhlo.broadcast_in_dim\"(%540) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3791 = \"mhlo.broadcast_in_dim\"(%565) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3792 = \"mhlo.broadcast_in_dim\"(%590) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3793 = \"mhlo.broadcast_in_dim\"(%615) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3794 = \"mhlo.broadcast_in_dim\"(%640) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3795 = \"mhlo.broadcast_in_dim\"(%665) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3796 = \"mhlo.broadcast_in_dim\"(%690) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3797 = \"mhlo.broadcast_in_dim\"(%715) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3798 = \"mhlo.broadcast_in_dim\"(%740) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3799 = \"mhlo.broadcast_in_dim\"(%765) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3800 = \"mhlo.broadcast_in_dim\"(%790) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3801 = \"mhlo.broadcast_in_dim\"(%815) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3802 = \"mhlo.broadcast_in_dim\"(%840) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3803 = \"mhlo.broadcast_in_dim\"(%865) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3804 = \"mhlo.broadcast_in_dim\"(%890) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3805 = \"mhlo.broadcast_in_dim\"(%915) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3806 = \"mhlo.broadcast_in_dim\"(%940) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3807 = \"mhlo.broadcast_in_dim\"(%965) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3808 = \"mhlo.broadcast_in_dim\"(%990) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3809 = \"mhlo.broadcast_in_dim\"(%1015) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3810 = \"mhlo.broadcast_in_dim\"(%1040) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3811 = \"mhlo.broadcast_in_dim\"(%1065) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3812 = \"mhlo.broadcast_in_dim\"(%1090) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3813 = \"mhlo.broadcast_in_dim\"(%1115) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3814 = \"mhlo.broadcast_in_dim\"(%1140) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3815 = \"mhlo.broadcast_in_dim\"(%1165) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3816 = \"mhlo.broadcast_in_dim\"(%1190) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3817 = \"mhlo.broadcast_in_dim\"(%1215) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3818 = \"mhlo.broadcast_in_dim\"(%1240) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3819 = \"mhlo.broadcast_in_dim\"(%1265) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3820 = \"mhlo.broadcast_in_dim\"(%1290) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3821 = \"mhlo.broadcast_in_dim\"(%1315) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3822 = \"mhlo.broadcast_in_dim\"(%1340) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3823 = \"mhlo.broadcast_in_dim\"(%1365) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3824 = \"mhlo.broadcast_in_dim\"(%1390) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3825 = \"mhlo.broadcast_in_dim\"(%1415) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3826 = \"mhlo.broadcast_in_dim\"(%1440) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3827 = \"mhlo.broadcast_in_dim\"(%1465) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3828 = \"mhlo.broadcast_in_dim\"(%1490) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3829 = \"mhlo.broadcast_in_dim\"(%1515) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3830 = \"mhlo.broadcast_in_dim\"(%1540) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3831 = \"mhlo.broadcast_in_dim\"(%1565) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3832 = \"mhlo.broadcast_in_dim\"(%1590) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3833 = \"mhlo.broadcast_in_dim\"(%1615) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3834 = \"mhlo.broadcast_in_dim\"(%1640) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3835 = \"mhlo.broadcast_in_dim\"(%1665) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3836 = \"mhlo.broadcast_in_dim\"(%1690) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3837 = \"mhlo.broadcast_in_dim\"(%1715) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3838 = \"mhlo.broadcast_in_dim\"(%1740) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3839 = \"mhlo.broadcast_in_dim\"(%1765) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3840 = \"mhlo.broadcast_in_dim\"(%1790) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3841 = \"mhlo.broadcast_in_dim\"(%1815) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3842 = \"mhlo.broadcast_in_dim\"(%1840) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3843 = \"mhlo.broadcast_in_dim\"(%1865) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3844 = \"mhlo.broadcast_in_dim\"(%1890) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3845 = \"mhlo.broadcast_in_dim\"(%1915) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3846 = \"mhlo.broadcast_in_dim\"(%1940) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3847 = \"mhlo.broadcast_in_dim\"(%1965) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3848 = \"mhlo.broadcast_in_dim\"(%1990) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3849 = \"mhlo.broadcast_in_dim\"(%2015) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3850 = \"mhlo.broadcast_in_dim\"(%2040) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3851 = \"mhlo.broadcast_in_dim\"(%2065) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3852 = \"mhlo.broadcast_in_dim\"(%2090) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3853 = \"mhlo.broadcast_in_dim\"(%2115) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3854 = \"mhlo.broadcast_in_dim\"(%2140) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3855 = \"mhlo.broadcast_in_dim\"(%2165) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3856 = \"mhlo.broadcast_in_dim\"(%2190) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3857 = \"mhlo.broadcast_in_dim\"(%2215) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3858 = \"mhlo.broadcast_in_dim\"(%2240) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3859 = \"mhlo.broadcast_in_dim\"(%2265) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3860 = \"mhlo.broadcast_in_dim\"(%2290) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3861 = \"mhlo.broadcast_in_dim\"(%2315) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3862 = \"mhlo.broadcast_in_dim\"(%2340) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3863 = \"mhlo.broadcast_in_dim\"(%2365) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3864 = \"mhlo.broadcast_in_dim\"(%2390) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3865 = \"mhlo.broadcast_in_dim\"(%2415) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3866 = \"mhlo.broadcast_in_dim\"(%2440) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3867 = \"mhlo.broadcast_in_dim\"(%2465) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3868 = \"mhlo.broadcast_in_dim\"(%2490) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3869 = \"mhlo.broadcast_in_dim\"(%2515) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3870 = \"mhlo.broadcast_in_dim\"(%2540) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3871 = \"mhlo.broadcast_in_dim\"(%2565) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3872 = \"mhlo.broadcast_in_dim\"(%2590) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3873 = \"mhlo.broadcast_in_dim\"(%2615) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3874 = \"mhlo.broadcast_in_dim\"(%2640) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3875 = \"mhlo.broadcast_in_dim\"(%2665) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3876 = \"mhlo.broadcast_in_dim\"(%2690) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3877 = \"mhlo.broadcast_in_dim\"(%2715) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3878 = \"mhlo.broadcast_in_dim\"(%2740) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3879 = \"mhlo.broadcast_in_dim\"(%2765) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3880 = \"mhlo.broadcast_in_dim\"(%2790) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3881 = \"mhlo.broadcast_in_dim\"(%2815) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3882 = \"mhlo.broadcast_in_dim\"(%2840) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3883 = \"mhlo.broadcast_in_dim\"(%2865) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3884 = \"mhlo.broadcast_in_dim\"(%2890) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3885 = \"mhlo.broadcast_in_dim\"(%2915) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3886 = \"mhlo.broadcast_in_dim\"(%2940) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3887 = \"mhlo.broadcast_in_dim\"(%2965) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3888 = \"mhlo.broadcast_in_dim\"(%2990) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3889 = \"mhlo.broadcast_in_dim\"(%3015) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3890 = \"mhlo.broadcast_in_dim\"(%3040) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3891 = \"mhlo.broadcast_in_dim\"(%3065) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3892 = \"mhlo.broadcast_in_dim\"(%3090) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3893 = \"mhlo.broadcast_in_dim\"(%3115) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3894 = \"mhlo.broadcast_in_dim\"(%3140) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3895 = \"mhlo.broadcast_in_dim\"(%3165) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3896 = \"mhlo.broadcast_in_dim\"(%3190) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3897 = \"mhlo.broadcast_in_dim\"(%3215) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3898 = \"mhlo.broadcast_in_dim\"(%3240) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3899 = \"mhlo.broadcast_in_dim\"(%3265) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3900 = \"mhlo.broadcast_in_dim\"(%3290) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3901 = \"mhlo.broadcast_in_dim\"(%3315) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3902 = \"mhlo.broadcast_in_dim\"(%3340) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3903 = \"mhlo.broadcast_in_dim\"(%3365) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3904 = \"mhlo.broadcast_in_dim\"(%3390) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3905 = \"mhlo.broadcast_in_dim\"(%3415) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3906 = \"mhlo.broadcast_in_dim\"(%3440) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3907 = \"mhlo.broadcast_in_dim\"(%3465) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3908 = \"mhlo.broadcast_in_dim\"(%3490) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3909 = \"mhlo.broadcast_in_dim\"(%3515) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3910 = \"mhlo.broadcast_in_dim\"(%3540) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3911 = \"mhlo.broadcast_in_dim\"(%3565) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3912 = \"mhlo.broadcast_in_dim\"(%3590) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3913 = \"mhlo.broadcast_in_dim\"(%3615) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<2x2xf32>) -> tensor<1x2x2xf32>\n",
      "    %3914 = \"mhlo.concatenate\"(%3770, %3771, %3772, %3773, %3774, %3775, %3776, %3777, %3778, %3779, %3780, %3781, %3782, %3783, %3784, %3785) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3915 = \"mhlo.concatenate\"(%3786, %3787, %3788, %3789, %3790, %3791, %3792, %3793, %3794, %3795, %3796, %3797, %3798, %3799, %3800, %3801) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3916 = \"mhlo.concatenate\"(%3802, %3803, %3804, %3805, %3806, %3807, %3808, %3809, %3810, %3811, %3812, %3813, %3814, %3815, %3816, %3817) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3917 = \"mhlo.concatenate\"(%3818, %3819, %3820, %3821, %3822, %3823, %3824, %3825, %3826, %3827, %3828, %3829, %3830, %3831, %3832, %3833) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3918 = \"mhlo.concatenate\"(%3834, %3835, %3836, %3837, %3838, %3839, %3840, %3841, %3842, %3843, %3844, %3845, %3846, %3847, %3848, %3849) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3919 = \"mhlo.concatenate\"(%3850, %3851, %3852, %3853, %3854, %3855, %3856, %3857, %3858, %3859, %3860, %3861, %3862, %3863, %3864, %3865) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3920 = \"mhlo.concatenate\"(%3866, %3867, %3868, %3869, %3870, %3871, %3872, %3873, %3874, %3875, %3876, %3877, %3878, %3879, %3880, %3881) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3921 = \"mhlo.concatenate\"(%3882, %3883, %3884, %3885, %3886, %3887, %3888, %3889, %3890, %3891, %3892, %3893, %3894, %3895, %3896, %3897) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3922 = \"mhlo.concatenate\"(%3898, %3899, %3900, %3901, %3902, %3903, %3904, %3905, %3906, %3907, %3908, %3909, %3910, %3911, %3912, %3913) {dimension = 0 : i64} : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor<1x2x2xf32>) -> tensor<16x2x2xf32>\n",
      "    %3923 = \"mhlo.concatenate\"(%3914, %3915, %3916, %3917, %3918, %3919, %3920, %3921, %3922) {dimension = 0 : i64} : (tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>, tensor<16x2x2xf32>) -> tensor<144x2x2xf32>\n",
      "    %3924 = \"mhlo.slice\"(%arg1) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<6x288xf32>) -> tensor<1x288xf32>\n",
      "    %3925 = mhlo.reshape %3924 : (tensor<1x288xf32>) -> tensor<288xf32>\n",
      "    %3926 = \"mhlo.dot_general\"(%15, %15) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288xf32>, tensor<288xf32>) -> tensor<f32>\n",
      "    %3927 = mhlo.constant dense<2.880000e+02> : tensor<f32>\n",
      "    %3928 = mhlo.divide %3926, %3927 : tensor<f32>\n",
      "    %3929 = mhlo.constant dense<9.99999974E-6> : tensor<f32>\n",
      "    %3930 = mhlo.add %3928, %3929 : tensor<f32>\n",
      "    %3931 = mhlo.sqrt %3930 : tensor<f32>\n",
      "    %3932 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %3933 = mhlo.divide %3932, %3931 : tensor<f32>\n",
      "    %3934 = mhlo.multiply %3925, %15 : tensor<288xf32>\n",
      "    %3935 = \"mhlo.broadcast_in_dim\"(%3933) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<288xf32>\n",
      "    %3936 = mhlo.multiply %3934, %3935 : tensor<288xf32>\n",
      "    %3937 = \"mhlo.slice\"(%arg11) {limit_indices = dense<[1, 288, 288]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %3938 = mhlo.reshape %3937 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %3939 = \"mhlo.dot_general\"(%3938, %3936) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %3940 = \"mhlo.slice\"(%arg9) {limit_indices = dense<[1, 288, 288]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %3941 = mhlo.reshape %3940 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %3942 = \"mhlo.dot_general\"(%3941, %3936) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %3943 = \"mhlo.slice\"(%arg12) {limit_indices = dense<[1, 288, 288]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %3944 = mhlo.reshape %3943 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %3945 = \"mhlo.dot_general\"(%3944, %3936) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %3946 = mhlo.reshape %3939 : (tensor<288xf32>) -> tensor<144x2xf32>\n",
      "    %3947 = mhlo.reshape %3942 : (tensor<288xf32>) -> tensor<144x2xf32>\n",
      "    %3948 = \"mhlo.dot_general\"(%3923, %3947) {dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>\n",
      "    %3949 = mhlo.reshape %3948 : (tensor<144x2xf32>) -> tensor<288xf32>\n",
      "    %3950 = \"mhlo.dot_general\"(%3769, %3946) {dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>\n",
      "    %3951 = mhlo.reshape %3950 : (tensor<144x2xf32>) -> tensor<288xf32>\n",
      "    %3952 = \"mhlo.slice\"(%arg13) {limit_indices = dense<[1, 0, 288]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x0x288xf32>) -> tensor<1x0x288xf32>\n",
      "    %3953 = mhlo.reshape %3952 : (tensor<1x0x288xf32>) -> tensor<0x288xf32>\n",
      "    %3954 = mhlo.reshape %3949 : (tensor<288xf32>) -> tensor<1x288xf32>\n",
      "    %3955 = call @append(%3953, %3954) : (tensor<0x288xf32>, tensor<1x288xf32>) -> tensor<1x288xf32>\n",
      "    %3956 = \"mhlo.slice\"(%arg14) {limit_indices = dense<[1, 0, 288]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x0x288xf32>) -> tensor<1x0x288xf32>\n",
      "    %3957 = mhlo.reshape %3956 : (tensor<1x0x288xf32>) -> tensor<0x288xf32>\n",
      "    %3958 = mhlo.reshape %3945 : (tensor<288xf32>) -> tensor<1x288xf32>\n",
      "    %3959 = call @append(%3957, %3958) : (tensor<0x288xf32>, tensor<1x288xf32>) -> tensor<1x288xf32>\n",
      "    %3960 = \"mhlo.slice\"(%3951) {limit_indices = dense<48> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %3961 = \"mhlo.slice\"(%3955) {limit_indices = dense<[1, 48]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %3962 = \"mhlo.dot_general\"(%3961, %3960) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %3963 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %3964 = mhlo.sqrt %3963 : tensor<f32>\n",
      "    %3965 = mhlo.convert %3964 : tensor<f32>\n",
      "    %3966 = \"mhlo.broadcast_in_dim\"(%3965) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3967 = mhlo.divide %3962, %3966 : tensor<1xf32>\n",
      "    %3968 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %3969 = mhlo.reduce(%3967 init: %3968) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %3970 = \"mhlo.broadcast_in_dim\"(%3969) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3971 = mhlo.subtract %3967, %3970 : tensor<1xf32>\n",
      "    %3972 = mhlo.exponential %3971 : tensor<1xf32>\n",
      "    %3973 = \"mhlo.slice\"(%3972) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %3974 = mhlo.reshape %3973 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %3975 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3976 = mhlo.add %3975, %3974 : tensor<f32>\n",
      "    %3977 = \"mhlo.broadcast_in_dim\"(%3976) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3978 = mhlo.divide %3972, %3977 : tensor<1xf32>\n",
      "    %3979 = \"mhlo.slice\"(%3959) {limit_indices = dense<[1, 48]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %3980 = \"mhlo.dot_general\"(%3979, %3978) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %3981 = \"mhlo.slice\"(%3951) {limit_indices = dense<96> : tensor<1xi64>, start_indices = dense<48> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %3982 = \"mhlo.slice\"(%3955) {limit_indices = dense<[1, 96]> : tensor<2xi64>, start_indices = dense<[0, 48]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %3983 = \"mhlo.dot_general\"(%3982, %3981) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %3984 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %3985 = mhlo.sqrt %3984 : tensor<f32>\n",
      "    %3986 = mhlo.convert %3985 : tensor<f32>\n",
      "    %3987 = \"mhlo.broadcast_in_dim\"(%3986) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3988 = mhlo.divide %3983, %3987 : tensor<1xf32>\n",
      "    %3989 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %3990 = mhlo.reduce(%3988 init: %3989) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %3991 = \"mhlo.broadcast_in_dim\"(%3990) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3992 = mhlo.subtract %3988, %3991 : tensor<1xf32>\n",
      "    %3993 = mhlo.exponential %3992 : tensor<1xf32>\n",
      "    %3994 = \"mhlo.slice\"(%3993) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %3995 = mhlo.reshape %3994 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %3996 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %3997 = mhlo.add %3996, %3995 : tensor<f32>\n",
      "    %3998 = \"mhlo.broadcast_in_dim\"(%3997) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %3999 = mhlo.divide %3993, %3998 : tensor<1xf32>\n",
      "    %4000 = \"mhlo.slice\"(%3959) {limit_indices = dense<[1, 96]> : tensor<2xi64>, start_indices = dense<[0, 48]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4001 = \"mhlo.dot_general\"(%4000, %3999) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4002 = \"mhlo.slice\"(%3951) {limit_indices = dense<144> : tensor<1xi64>, start_indices = dense<96> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4003 = \"mhlo.slice\"(%3955) {limit_indices = dense<[1, 144]> : tensor<2xi64>, start_indices = dense<[0, 96]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4004 = \"mhlo.dot_general\"(%4003, %4002) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4005 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4006 = mhlo.sqrt %4005 : tensor<f32>\n",
      "    %4007 = mhlo.convert %4006 : tensor<f32>\n",
      "    %4008 = \"mhlo.broadcast_in_dim\"(%4007) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4009 = mhlo.divide %4004, %4008 : tensor<1xf32>\n",
      "    %4010 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4011 = mhlo.reduce(%4009 init: %4010) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4012 = \"mhlo.broadcast_in_dim\"(%4011) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4013 = mhlo.subtract %4009, %4012 : tensor<1xf32>\n",
      "    %4014 = mhlo.exponential %4013 : tensor<1xf32>\n",
      "    %4015 = \"mhlo.slice\"(%4014) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4016 = mhlo.reshape %4015 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4017 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4018 = mhlo.add %4017, %4016 : tensor<f32>\n",
      "    %4019 = \"mhlo.broadcast_in_dim\"(%4018) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4020 = mhlo.divide %4014, %4019 : tensor<1xf32>\n",
      "    %4021 = \"mhlo.slice\"(%3959) {limit_indices = dense<[1, 144]> : tensor<2xi64>, start_indices = dense<[0, 96]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4022 = \"mhlo.dot_general\"(%4021, %4020) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4023 = \"mhlo.slice\"(%3951) {limit_indices = dense<192> : tensor<1xi64>, start_indices = dense<144> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4024 = \"mhlo.slice\"(%3955) {limit_indices = dense<[1, 192]> : tensor<2xi64>, start_indices = dense<[0, 144]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4025 = \"mhlo.dot_general\"(%4024, %4023) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4026 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4027 = mhlo.sqrt %4026 : tensor<f32>\n",
      "    %4028 = mhlo.convert %4027 : tensor<f32>\n",
      "    %4029 = \"mhlo.broadcast_in_dim\"(%4028) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4030 = mhlo.divide %4025, %4029 : tensor<1xf32>\n",
      "    %4031 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4032 = mhlo.reduce(%4030 init: %4031) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4033 = \"mhlo.broadcast_in_dim\"(%4032) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4034 = mhlo.subtract %4030, %4033 : tensor<1xf32>\n",
      "    %4035 = mhlo.exponential %4034 : tensor<1xf32>\n",
      "    %4036 = \"mhlo.slice\"(%4035) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4037 = mhlo.reshape %4036 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4038 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4039 = mhlo.add %4038, %4037 : tensor<f32>\n",
      "    %4040 = \"mhlo.broadcast_in_dim\"(%4039) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4041 = mhlo.divide %4035, %4040 : tensor<1xf32>\n",
      "    %4042 = \"mhlo.slice\"(%3959) {limit_indices = dense<[1, 192]> : tensor<2xi64>, start_indices = dense<[0, 144]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4043 = \"mhlo.dot_general\"(%4042, %4041) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4044 = \"mhlo.slice\"(%3951) {limit_indices = dense<240> : tensor<1xi64>, start_indices = dense<192> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4045 = \"mhlo.slice\"(%3955) {limit_indices = dense<[1, 240]> : tensor<2xi64>, start_indices = dense<[0, 192]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4046 = \"mhlo.dot_general\"(%4045, %4044) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4047 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4048 = mhlo.sqrt %4047 : tensor<f32>\n",
      "    %4049 = mhlo.convert %4048 : tensor<f32>\n",
      "    %4050 = \"mhlo.broadcast_in_dim\"(%4049) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4051 = mhlo.divide %4046, %4050 : tensor<1xf32>\n",
      "    %4052 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4053 = mhlo.reduce(%4051 init: %4052) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4054 = \"mhlo.broadcast_in_dim\"(%4053) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4055 = mhlo.subtract %4051, %4054 : tensor<1xf32>\n",
      "    %4056 = mhlo.exponential %4055 : tensor<1xf32>\n",
      "    %4057 = \"mhlo.slice\"(%4056) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4058 = mhlo.reshape %4057 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4059 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4060 = mhlo.add %4059, %4058 : tensor<f32>\n",
      "    %4061 = \"mhlo.broadcast_in_dim\"(%4060) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4062 = mhlo.divide %4056, %4061 : tensor<1xf32>\n",
      "    %4063 = \"mhlo.slice\"(%3959) {limit_indices = dense<[1, 240]> : tensor<2xi64>, start_indices = dense<[0, 192]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4064 = \"mhlo.dot_general\"(%4063, %4062) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4065 = \"mhlo.slice\"(%3951) {limit_indices = dense<288> : tensor<1xi64>, start_indices = dense<240> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4066 = \"mhlo.slice\"(%3955) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<[0, 240]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4067 = \"mhlo.dot_general\"(%4066, %4065) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4068 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4069 = mhlo.sqrt %4068 : tensor<f32>\n",
      "    %4070 = mhlo.convert %4069 : tensor<f32>\n",
      "    %4071 = \"mhlo.broadcast_in_dim\"(%4070) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4072 = mhlo.divide %4067, %4071 : tensor<1xf32>\n",
      "    %4073 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4074 = mhlo.reduce(%4072 init: %4073) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4075 = \"mhlo.broadcast_in_dim\"(%4074) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4076 = mhlo.subtract %4072, %4075 : tensor<1xf32>\n",
      "    %4077 = mhlo.exponential %4076 : tensor<1xf32>\n",
      "    %4078 = \"mhlo.slice\"(%4077) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4079 = mhlo.reshape %4078 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4080 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4081 = mhlo.add %4080, %4079 : tensor<f32>\n",
      "    %4082 = \"mhlo.broadcast_in_dim\"(%4081) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4083 = mhlo.divide %4077, %4082 : tensor<1xf32>\n",
      "    %4084 = \"mhlo.slice\"(%3959) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<[0, 240]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4085 = \"mhlo.dot_general\"(%4084, %4083) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4086 = \"mhlo.concatenate\"(%3980, %4001, %4022, %4043, %4064, %4085) {dimension = 0 : i64} : (tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>) -> tensor<288xf32>\n",
      "    %4087 = \"mhlo.slice\"(%arg10) {limit_indices = dense<[1, 288, 288]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4088 = mhlo.reshape %4087 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4089 = \"mhlo.dot_general\"(%4088, %4086) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4090 = mhlo.add %15, %4089 : tensor<288xf32>\n",
      "    %4091 = \"mhlo.slice\"(%arg2) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<6x288xf32>) -> tensor<1x288xf32>\n",
      "    %4092 = mhlo.reshape %4091 : (tensor<1x288xf32>) -> tensor<288xf32>\n",
      "    %4093 = \"mhlo.dot_general\"(%4090, %4090) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288xf32>, tensor<288xf32>) -> tensor<f32>\n",
      "    %4094 = mhlo.constant dense<2.880000e+02> : tensor<f32>\n",
      "    %4095 = mhlo.divide %4093, %4094 : tensor<f32>\n",
      "    %4096 = mhlo.constant dense<9.99999974E-6> : tensor<f32>\n",
      "    %4097 = mhlo.add %4095, %4096 : tensor<f32>\n",
      "    %4098 = mhlo.sqrt %4097 : tensor<f32>\n",
      "    %4099 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4100 = mhlo.divide %4099, %4098 : tensor<f32>\n",
      "    %4101 = mhlo.multiply %4092, %4090 : tensor<288xf32>\n",
      "    %4102 = \"mhlo.broadcast_in_dim\"(%4100) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<288xf32>\n",
      "    %4103 = mhlo.multiply %4101, %4102 : tensor<288xf32>\n",
      "    %4104 = \"mhlo.slice\"(%arg5) {limit_indices = dense<[1, 768, 288]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x768x288xf32>) -> tensor<1x768x288xf32>\n",
      "    %4105 = mhlo.reshape %4104 : (tensor<1x768x288xf32>) -> tensor<768x288xf32>\n",
      "    %4106 = \"mhlo.dot_general\"(%4105, %4103) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<768x288xf32>, tensor<288xf32>) -> tensor<768xf32>\n",
      "    %4107 = \"mhlo.slice\"(%arg7) {limit_indices = dense<[1, 768, 288]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x768x288xf32>) -> tensor<1x768x288xf32>\n",
      "    %4108 = mhlo.reshape %4107 : (tensor<1x768x288xf32>) -> tensor<768x288xf32>\n",
      "    %4109 = \"mhlo.dot_general\"(%4108, %4103) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<768x288xf32>, tensor<288xf32>) -> tensor<768xf32>\n",
      "    %4110 = mhlo.negate %4106 : tensor<768xf32>\n",
      "    %4111 = mhlo.exponential %4110 : tensor<768xf32>\n",
      "    %4112 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4113 = \"mhlo.broadcast_in_dim\"(%4112) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<768xf32>\n",
      "    %4114 = mhlo.add %4113, %4111 : tensor<768xf32>\n",
      "    %4115 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4116 = \"mhlo.broadcast_in_dim\"(%4115) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<768xf32>\n",
      "    %4117 = mhlo.divide %4116, %4114 : tensor<768xf32>\n",
      "    %4118 = mhlo.multiply %4106, %4117 : tensor<768xf32>\n",
      "    %4119 = mhlo.multiply %4118, %4109 : tensor<768xf32>\n",
      "    %4120 = \"mhlo.slice\"(%arg6) {limit_indices = dense<[1, 288, 768]> : tensor<3xi64>, start_indices = dense<0> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x768xf32>) -> tensor<1x288x768xf32>\n",
      "    %4121 = mhlo.reshape %4120 : (tensor<1x288x768xf32>) -> tensor<288x768xf32>\n",
      "    %4122 = \"mhlo.dot_general\"(%4121, %4119) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x768xf32>, tensor<768xf32>) -> tensor<288xf32>\n",
      "    %4123 = mhlo.add %4090, %4122 : tensor<288xf32>\n",
      "    %4124 = \"mhlo.slice\"(%arg1) {limit_indices = dense<[2, 288]> : tensor<2xi64>, start_indices = dense<[1, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<6x288xf32>) -> tensor<1x288xf32>\n",
      "    %4125 = mhlo.reshape %4124 : (tensor<1x288xf32>) -> tensor<288xf32>\n",
      "    %4126 = \"mhlo.dot_general\"(%4123, %4123) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288xf32>, tensor<288xf32>) -> tensor<f32>\n",
      "    %4127 = mhlo.constant dense<2.880000e+02> : tensor<f32>\n",
      "    %4128 = mhlo.divide %4126, %4127 : tensor<f32>\n",
      "    %4129 = mhlo.constant dense<9.99999974E-6> : tensor<f32>\n",
      "    %4130 = mhlo.add %4128, %4129 : tensor<f32>\n",
      "    %4131 = mhlo.sqrt %4130 : tensor<f32>\n",
      "    %4132 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4133 = mhlo.divide %4132, %4131 : tensor<f32>\n",
      "    %4134 = mhlo.multiply %4125, %4123 : tensor<288xf32>\n",
      "    %4135 = \"mhlo.broadcast_in_dim\"(%4133) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<288xf32>\n",
      "    %4136 = mhlo.multiply %4134, %4135 : tensor<288xf32>\n",
      "    %4137 = \"mhlo.slice\"(%arg11) {limit_indices = dense<[2, 288, 288]> : tensor<3xi64>, start_indices = dense<[1, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4138 = mhlo.reshape %4137 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4139 = \"mhlo.dot_general\"(%4138, %4136) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4140 = \"mhlo.slice\"(%arg9) {limit_indices = dense<[2, 288, 288]> : tensor<3xi64>, start_indices = dense<[1, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4141 = mhlo.reshape %4140 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4142 = \"mhlo.dot_general\"(%4141, %4136) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4143 = \"mhlo.slice\"(%arg12) {limit_indices = dense<[2, 288, 288]> : tensor<3xi64>, start_indices = dense<[1, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4144 = mhlo.reshape %4143 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4145 = \"mhlo.dot_general\"(%4144, %4136) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4146 = mhlo.reshape %4139 : (tensor<288xf32>) -> tensor<144x2xf32>\n",
      "    %4147 = mhlo.reshape %4142 : (tensor<288xf32>) -> tensor<144x2xf32>\n",
      "    %4148 = \"mhlo.dot_general\"(%3923, %4147) {dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>\n",
      "    %4149 = mhlo.reshape %4148 : (tensor<144x2xf32>) -> tensor<288xf32>\n",
      "    %4150 = \"mhlo.dot_general\"(%3769, %4146) {dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>\n",
      "    %4151 = mhlo.reshape %4150 : (tensor<144x2xf32>) -> tensor<288xf32>\n",
      "    %4152 = \"mhlo.slice\"(%arg13) {limit_indices = dense<[2, 0, 288]> : tensor<3xi64>, start_indices = dense<[1, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x0x288xf32>) -> tensor<1x0x288xf32>\n",
      "    %4153 = mhlo.reshape %4152 : (tensor<1x0x288xf32>) -> tensor<0x288xf32>\n",
      "    %4154 = mhlo.reshape %4149 : (tensor<288xf32>) -> tensor<1x288xf32>\n",
      "    %4155 = call @append(%4153, %4154) : (tensor<0x288xf32>, tensor<1x288xf32>) -> tensor<1x288xf32>\n",
      "    %4156 = \"mhlo.slice\"(%arg14) {limit_indices = dense<[2, 0, 288]> : tensor<3xi64>, start_indices = dense<[1, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x0x288xf32>) -> tensor<1x0x288xf32>\n",
      "    %4157 = mhlo.reshape %4156 : (tensor<1x0x288xf32>) -> tensor<0x288xf32>\n",
      "    %4158 = mhlo.reshape %4145 : (tensor<288xf32>) -> tensor<1x288xf32>\n",
      "    %4159 = call @append(%4157, %4158) : (tensor<0x288xf32>, tensor<1x288xf32>) -> tensor<1x288xf32>\n",
      "    %4160 = \"mhlo.slice\"(%4151) {limit_indices = dense<48> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4161 = \"mhlo.slice\"(%4155) {limit_indices = dense<[1, 48]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4162 = \"mhlo.dot_general\"(%4161, %4160) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4163 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4164 = mhlo.sqrt %4163 : tensor<f32>\n",
      "    %4165 = mhlo.convert %4164 : tensor<f32>\n",
      "    %4166 = \"mhlo.broadcast_in_dim\"(%4165) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4167 = mhlo.divide %4162, %4166 : tensor<1xf32>\n",
      "    %4168 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4169 = mhlo.reduce(%4167 init: %4168) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4170 = \"mhlo.broadcast_in_dim\"(%4169) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4171 = mhlo.subtract %4167, %4170 : tensor<1xf32>\n",
      "    %4172 = mhlo.exponential %4171 : tensor<1xf32>\n",
      "    %4173 = \"mhlo.slice\"(%4172) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4174 = mhlo.reshape %4173 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4175 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4176 = mhlo.add %4175, %4174 : tensor<f32>\n",
      "    %4177 = \"mhlo.broadcast_in_dim\"(%4176) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4178 = mhlo.divide %4172, %4177 : tensor<1xf32>\n",
      "    %4179 = \"mhlo.slice\"(%4159) {limit_indices = dense<[1, 48]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4180 = \"mhlo.dot_general\"(%4179, %4178) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4181 = \"mhlo.slice\"(%4151) {limit_indices = dense<96> : tensor<1xi64>, start_indices = dense<48> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4182 = \"mhlo.slice\"(%4155) {limit_indices = dense<[1, 96]> : tensor<2xi64>, start_indices = dense<[0, 48]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4183 = \"mhlo.dot_general\"(%4182, %4181) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4184 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4185 = mhlo.sqrt %4184 : tensor<f32>\n",
      "    %4186 = mhlo.convert %4185 : tensor<f32>\n",
      "    %4187 = \"mhlo.broadcast_in_dim\"(%4186) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4188 = mhlo.divide %4183, %4187 : tensor<1xf32>\n",
      "    %4189 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4190 = mhlo.reduce(%4188 init: %4189) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4191 = \"mhlo.broadcast_in_dim\"(%4190) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4192 = mhlo.subtract %4188, %4191 : tensor<1xf32>\n",
      "    %4193 = mhlo.exponential %4192 : tensor<1xf32>\n",
      "    %4194 = \"mhlo.slice\"(%4193) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4195 = mhlo.reshape %4194 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4196 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4197 = mhlo.add %4196, %4195 : tensor<f32>\n",
      "    %4198 = \"mhlo.broadcast_in_dim\"(%4197) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4199 = mhlo.divide %4193, %4198 : tensor<1xf32>\n",
      "    %4200 = \"mhlo.slice\"(%4159) {limit_indices = dense<[1, 96]> : tensor<2xi64>, start_indices = dense<[0, 48]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4201 = \"mhlo.dot_general\"(%4200, %4199) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4202 = \"mhlo.slice\"(%4151) {limit_indices = dense<144> : tensor<1xi64>, start_indices = dense<96> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4203 = \"mhlo.slice\"(%4155) {limit_indices = dense<[1, 144]> : tensor<2xi64>, start_indices = dense<[0, 96]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4204 = \"mhlo.dot_general\"(%4203, %4202) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4205 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4206 = mhlo.sqrt %4205 : tensor<f32>\n",
      "    %4207 = mhlo.convert %4206 : tensor<f32>\n",
      "    %4208 = \"mhlo.broadcast_in_dim\"(%4207) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4209 = mhlo.divide %4204, %4208 : tensor<1xf32>\n",
      "    %4210 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4211 = mhlo.reduce(%4209 init: %4210) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4212 = \"mhlo.broadcast_in_dim\"(%4211) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4213 = mhlo.subtract %4209, %4212 : tensor<1xf32>\n",
      "    %4214 = mhlo.exponential %4213 : tensor<1xf32>\n",
      "    %4215 = \"mhlo.slice\"(%4214) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4216 = mhlo.reshape %4215 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4217 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4218 = mhlo.add %4217, %4216 : tensor<f32>\n",
      "    %4219 = \"mhlo.broadcast_in_dim\"(%4218) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4220 = mhlo.divide %4214, %4219 : tensor<1xf32>\n",
      "    %4221 = \"mhlo.slice\"(%4159) {limit_indices = dense<[1, 144]> : tensor<2xi64>, start_indices = dense<[0, 96]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4222 = \"mhlo.dot_general\"(%4221, %4220) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4223 = \"mhlo.slice\"(%4151) {limit_indices = dense<192> : tensor<1xi64>, start_indices = dense<144> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4224 = \"mhlo.slice\"(%4155) {limit_indices = dense<[1, 192]> : tensor<2xi64>, start_indices = dense<[0, 144]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4225 = \"mhlo.dot_general\"(%4224, %4223) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4226 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4227 = mhlo.sqrt %4226 : tensor<f32>\n",
      "    %4228 = mhlo.convert %4227 : tensor<f32>\n",
      "    %4229 = \"mhlo.broadcast_in_dim\"(%4228) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4230 = mhlo.divide %4225, %4229 : tensor<1xf32>\n",
      "    %4231 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4232 = mhlo.reduce(%4230 init: %4231) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4233 = \"mhlo.broadcast_in_dim\"(%4232) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4234 = mhlo.subtract %4230, %4233 : tensor<1xf32>\n",
      "    %4235 = mhlo.exponential %4234 : tensor<1xf32>\n",
      "    %4236 = \"mhlo.slice\"(%4235) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4237 = mhlo.reshape %4236 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4238 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4239 = mhlo.add %4238, %4237 : tensor<f32>\n",
      "    %4240 = \"mhlo.broadcast_in_dim\"(%4239) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4241 = mhlo.divide %4235, %4240 : tensor<1xf32>\n",
      "    %4242 = \"mhlo.slice\"(%4159) {limit_indices = dense<[1, 192]> : tensor<2xi64>, start_indices = dense<[0, 144]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4243 = \"mhlo.dot_general\"(%4242, %4241) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4244 = \"mhlo.slice\"(%4151) {limit_indices = dense<240> : tensor<1xi64>, start_indices = dense<192> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4245 = \"mhlo.slice\"(%4155) {limit_indices = dense<[1, 240]> : tensor<2xi64>, start_indices = dense<[0, 192]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4246 = \"mhlo.dot_general\"(%4245, %4244) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4247 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4248 = mhlo.sqrt %4247 : tensor<f32>\n",
      "    %4249 = mhlo.convert %4248 : tensor<f32>\n",
      "    %4250 = \"mhlo.broadcast_in_dim\"(%4249) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4251 = mhlo.divide %4246, %4250 : tensor<1xf32>\n",
      "    %4252 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4253 = mhlo.reduce(%4251 init: %4252) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4254 = \"mhlo.broadcast_in_dim\"(%4253) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4255 = mhlo.subtract %4251, %4254 : tensor<1xf32>\n",
      "    %4256 = mhlo.exponential %4255 : tensor<1xf32>\n",
      "    %4257 = \"mhlo.slice\"(%4256) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4258 = mhlo.reshape %4257 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4259 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4260 = mhlo.add %4259, %4258 : tensor<f32>\n",
      "    %4261 = \"mhlo.broadcast_in_dim\"(%4260) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4262 = mhlo.divide %4256, %4261 : tensor<1xf32>\n",
      "    %4263 = \"mhlo.slice\"(%4159) {limit_indices = dense<[1, 240]> : tensor<2xi64>, start_indices = dense<[0, 192]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4264 = \"mhlo.dot_general\"(%4263, %4262) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4265 = \"mhlo.slice\"(%4151) {limit_indices = dense<288> : tensor<1xi64>, start_indices = dense<240> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4266 = \"mhlo.slice\"(%4155) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<[0, 240]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4267 = \"mhlo.dot_general\"(%4266, %4265) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4268 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4269 = mhlo.sqrt %4268 : tensor<f32>\n",
      "    %4270 = mhlo.convert %4269 : tensor<f32>\n",
      "    %4271 = \"mhlo.broadcast_in_dim\"(%4270) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4272 = mhlo.divide %4267, %4271 : tensor<1xf32>\n",
      "    %4273 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4274 = mhlo.reduce(%4272 init: %4273) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4275 = \"mhlo.broadcast_in_dim\"(%4274) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4276 = mhlo.subtract %4272, %4275 : tensor<1xf32>\n",
      "    %4277 = mhlo.exponential %4276 : tensor<1xf32>\n",
      "    %4278 = \"mhlo.slice\"(%4277) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4279 = mhlo.reshape %4278 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4280 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4281 = mhlo.add %4280, %4279 : tensor<f32>\n",
      "    %4282 = \"mhlo.broadcast_in_dim\"(%4281) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4283 = mhlo.divide %4277, %4282 : tensor<1xf32>\n",
      "    %4284 = \"mhlo.slice\"(%4159) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<[0, 240]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4285 = \"mhlo.dot_general\"(%4284, %4283) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4286 = \"mhlo.concatenate\"(%4180, %4201, %4222, %4243, %4264, %4285) {dimension = 0 : i64} : (tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>) -> tensor<288xf32>\n",
      "    %4287 = \"mhlo.slice\"(%arg10) {limit_indices = dense<[2, 288, 288]> : tensor<3xi64>, start_indices = dense<[1, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4288 = mhlo.reshape %4287 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4289 = \"mhlo.dot_general\"(%4288, %4286) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4290 = mhlo.add %4123, %4289 : tensor<288xf32>\n",
      "    %4291 = \"mhlo.slice\"(%arg2) {limit_indices = dense<[2, 288]> : tensor<2xi64>, start_indices = dense<[1, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<6x288xf32>) -> tensor<1x288xf32>\n",
      "    %4292 = mhlo.reshape %4291 : (tensor<1x288xf32>) -> tensor<288xf32>\n",
      "    %4293 = \"mhlo.dot_general\"(%4290, %4290) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288xf32>, tensor<288xf32>) -> tensor<f32>\n",
      "    %4294 = mhlo.constant dense<2.880000e+02> : tensor<f32>\n",
      "    %4295 = mhlo.divide %4293, %4294 : tensor<f32>\n",
      "    %4296 = mhlo.constant dense<9.99999974E-6> : tensor<f32>\n",
      "    %4297 = mhlo.add %4295, %4296 : tensor<f32>\n",
      "    %4298 = mhlo.sqrt %4297 : tensor<f32>\n",
      "    %4299 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4300 = mhlo.divide %4299, %4298 : tensor<f32>\n",
      "    %4301 = mhlo.multiply %4292, %4290 : tensor<288xf32>\n",
      "    %4302 = \"mhlo.broadcast_in_dim\"(%4300) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<288xf32>\n",
      "    %4303 = mhlo.multiply %4301, %4302 : tensor<288xf32>\n",
      "    %4304 = \"mhlo.slice\"(%arg5) {limit_indices = dense<[2, 768, 288]> : tensor<3xi64>, start_indices = dense<[1, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x768x288xf32>) -> tensor<1x768x288xf32>\n",
      "    %4305 = mhlo.reshape %4304 : (tensor<1x768x288xf32>) -> tensor<768x288xf32>\n",
      "    %4306 = \"mhlo.dot_general\"(%4305, %4303) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<768x288xf32>, tensor<288xf32>) -> tensor<768xf32>\n",
      "    %4307 = \"mhlo.slice\"(%arg7) {limit_indices = dense<[2, 768, 288]> : tensor<3xi64>, start_indices = dense<[1, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x768x288xf32>) -> tensor<1x768x288xf32>\n",
      "    %4308 = mhlo.reshape %4307 : (tensor<1x768x288xf32>) -> tensor<768x288xf32>\n",
      "    %4309 = \"mhlo.dot_general\"(%4308, %4303) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<768x288xf32>, tensor<288xf32>) -> tensor<768xf32>\n",
      "    %4310 = mhlo.negate %4306 : tensor<768xf32>\n",
      "    %4311 = mhlo.exponential %4310 : tensor<768xf32>\n",
      "    %4312 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4313 = \"mhlo.broadcast_in_dim\"(%4312) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<768xf32>\n",
      "    %4314 = mhlo.add %4313, %4311 : tensor<768xf32>\n",
      "    %4315 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4316 = \"mhlo.broadcast_in_dim\"(%4315) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<768xf32>\n",
      "    %4317 = mhlo.divide %4316, %4314 : tensor<768xf32>\n",
      "    %4318 = mhlo.multiply %4306, %4317 : tensor<768xf32>\n",
      "    %4319 = mhlo.multiply %4318, %4309 : tensor<768xf32>\n",
      "    %4320 = \"mhlo.slice\"(%arg6) {limit_indices = dense<[2, 288, 768]> : tensor<3xi64>, start_indices = dense<[1, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x768xf32>) -> tensor<1x288x768xf32>\n",
      "    %4321 = mhlo.reshape %4320 : (tensor<1x288x768xf32>) -> tensor<288x768xf32>\n",
      "    %4322 = \"mhlo.dot_general\"(%4321, %4319) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x768xf32>, tensor<768xf32>) -> tensor<288xf32>\n",
      "    %4323 = mhlo.add %4290, %4322 : tensor<288xf32>\n",
      "    %4324 = \"mhlo.slice\"(%arg1) {limit_indices = dense<[3, 288]> : tensor<2xi64>, start_indices = dense<[2, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<6x288xf32>) -> tensor<1x288xf32>\n",
      "    %4325 = mhlo.reshape %4324 : (tensor<1x288xf32>) -> tensor<288xf32>\n",
      "    %4326 = \"mhlo.dot_general\"(%4323, %4323) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288xf32>, tensor<288xf32>) -> tensor<f32>\n",
      "    %4327 = mhlo.constant dense<2.880000e+02> : tensor<f32>\n",
      "    %4328 = mhlo.divide %4326, %4327 : tensor<f32>\n",
      "    %4329 = mhlo.constant dense<9.99999974E-6> : tensor<f32>\n",
      "    %4330 = mhlo.add %4328, %4329 : tensor<f32>\n",
      "    %4331 = mhlo.sqrt %4330 : tensor<f32>\n",
      "    %4332 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4333 = mhlo.divide %4332, %4331 : tensor<f32>\n",
      "    %4334 = mhlo.multiply %4325, %4323 : tensor<288xf32>\n",
      "    %4335 = \"mhlo.broadcast_in_dim\"(%4333) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<288xf32>\n",
      "    %4336 = mhlo.multiply %4334, %4335 : tensor<288xf32>\n",
      "    %4337 = \"mhlo.slice\"(%arg11) {limit_indices = dense<[3, 288, 288]> : tensor<3xi64>, start_indices = dense<[2, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4338 = mhlo.reshape %4337 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4339 = \"mhlo.dot_general\"(%4338, %4336) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4340 = \"mhlo.slice\"(%arg9) {limit_indices = dense<[3, 288, 288]> : tensor<3xi64>, start_indices = dense<[2, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4341 = mhlo.reshape %4340 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4342 = \"mhlo.dot_general\"(%4341, %4336) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4343 = \"mhlo.slice\"(%arg12) {limit_indices = dense<[3, 288, 288]> : tensor<3xi64>, start_indices = dense<[2, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4344 = mhlo.reshape %4343 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4345 = \"mhlo.dot_general\"(%4344, %4336) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4346 = mhlo.reshape %4339 : (tensor<288xf32>) -> tensor<144x2xf32>\n",
      "    %4347 = mhlo.reshape %4342 : (tensor<288xf32>) -> tensor<144x2xf32>\n",
      "    %4348 = \"mhlo.dot_general\"(%3923, %4347) {dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>\n",
      "    %4349 = mhlo.reshape %4348 : (tensor<144x2xf32>) -> tensor<288xf32>\n",
      "    %4350 = \"mhlo.dot_general\"(%3769, %4346) {dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>\n",
      "    %4351 = mhlo.reshape %4350 : (tensor<144x2xf32>) -> tensor<288xf32>\n",
      "    %4352 = \"mhlo.slice\"(%arg13) {limit_indices = dense<[3, 0, 288]> : tensor<3xi64>, start_indices = dense<[2, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x0x288xf32>) -> tensor<1x0x288xf32>\n",
      "    %4353 = mhlo.reshape %4352 : (tensor<1x0x288xf32>) -> tensor<0x288xf32>\n",
      "    %4354 = mhlo.reshape %4349 : (tensor<288xf32>) -> tensor<1x288xf32>\n",
      "    %4355 = call @append(%4353, %4354) : (tensor<0x288xf32>, tensor<1x288xf32>) -> tensor<1x288xf32>\n",
      "    %4356 = \"mhlo.slice\"(%arg14) {limit_indices = dense<[3, 0, 288]> : tensor<3xi64>, start_indices = dense<[2, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x0x288xf32>) -> tensor<1x0x288xf32>\n",
      "    %4357 = mhlo.reshape %4356 : (tensor<1x0x288xf32>) -> tensor<0x288xf32>\n",
      "    %4358 = mhlo.reshape %4345 : (tensor<288xf32>) -> tensor<1x288xf32>\n",
      "    %4359 = call @append(%4357, %4358) : (tensor<0x288xf32>, tensor<1x288xf32>) -> tensor<1x288xf32>\n",
      "    %4360 = \"mhlo.slice\"(%4351) {limit_indices = dense<48> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4361 = \"mhlo.slice\"(%4355) {limit_indices = dense<[1, 48]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4362 = \"mhlo.dot_general\"(%4361, %4360) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4363 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4364 = mhlo.sqrt %4363 : tensor<f32>\n",
      "    %4365 = mhlo.convert %4364 : tensor<f32>\n",
      "    %4366 = \"mhlo.broadcast_in_dim\"(%4365) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4367 = mhlo.divide %4362, %4366 : tensor<1xf32>\n",
      "    %4368 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4369 = mhlo.reduce(%4367 init: %4368) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4370 = \"mhlo.broadcast_in_dim\"(%4369) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4371 = mhlo.subtract %4367, %4370 : tensor<1xf32>\n",
      "    %4372 = mhlo.exponential %4371 : tensor<1xf32>\n",
      "    %4373 = \"mhlo.slice\"(%4372) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4374 = mhlo.reshape %4373 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4375 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4376 = mhlo.add %4375, %4374 : tensor<f32>\n",
      "    %4377 = \"mhlo.broadcast_in_dim\"(%4376) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4378 = mhlo.divide %4372, %4377 : tensor<1xf32>\n",
      "    %4379 = \"mhlo.slice\"(%4359) {limit_indices = dense<[1, 48]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4380 = \"mhlo.dot_general\"(%4379, %4378) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4381 = \"mhlo.slice\"(%4351) {limit_indices = dense<96> : tensor<1xi64>, start_indices = dense<48> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4382 = \"mhlo.slice\"(%4355) {limit_indices = dense<[1, 96]> : tensor<2xi64>, start_indices = dense<[0, 48]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4383 = \"mhlo.dot_general\"(%4382, %4381) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4384 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4385 = mhlo.sqrt %4384 : tensor<f32>\n",
      "    %4386 = mhlo.convert %4385 : tensor<f32>\n",
      "    %4387 = \"mhlo.broadcast_in_dim\"(%4386) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4388 = mhlo.divide %4383, %4387 : tensor<1xf32>\n",
      "    %4389 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4390 = mhlo.reduce(%4388 init: %4389) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4391 = \"mhlo.broadcast_in_dim\"(%4390) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4392 = mhlo.subtract %4388, %4391 : tensor<1xf32>\n",
      "    %4393 = mhlo.exponential %4392 : tensor<1xf32>\n",
      "    %4394 = \"mhlo.slice\"(%4393) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4395 = mhlo.reshape %4394 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4396 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4397 = mhlo.add %4396, %4395 : tensor<f32>\n",
      "    %4398 = \"mhlo.broadcast_in_dim\"(%4397) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4399 = mhlo.divide %4393, %4398 : tensor<1xf32>\n",
      "    %4400 = \"mhlo.slice\"(%4359) {limit_indices = dense<[1, 96]> : tensor<2xi64>, start_indices = dense<[0, 48]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4401 = \"mhlo.dot_general\"(%4400, %4399) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4402 = \"mhlo.slice\"(%4351) {limit_indices = dense<144> : tensor<1xi64>, start_indices = dense<96> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4403 = \"mhlo.slice\"(%4355) {limit_indices = dense<[1, 144]> : tensor<2xi64>, start_indices = dense<[0, 96]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4404 = \"mhlo.dot_general\"(%4403, %4402) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4405 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4406 = mhlo.sqrt %4405 : tensor<f32>\n",
      "    %4407 = mhlo.convert %4406 : tensor<f32>\n",
      "    %4408 = \"mhlo.broadcast_in_dim\"(%4407) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4409 = mhlo.divide %4404, %4408 : tensor<1xf32>\n",
      "    %4410 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4411 = mhlo.reduce(%4409 init: %4410) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4412 = \"mhlo.broadcast_in_dim\"(%4411) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4413 = mhlo.subtract %4409, %4412 : tensor<1xf32>\n",
      "    %4414 = mhlo.exponential %4413 : tensor<1xf32>\n",
      "    %4415 = \"mhlo.slice\"(%4414) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4416 = mhlo.reshape %4415 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4417 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4418 = mhlo.add %4417, %4416 : tensor<f32>\n",
      "    %4419 = \"mhlo.broadcast_in_dim\"(%4418) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4420 = mhlo.divide %4414, %4419 : tensor<1xf32>\n",
      "    %4421 = \"mhlo.slice\"(%4359) {limit_indices = dense<[1, 144]> : tensor<2xi64>, start_indices = dense<[0, 96]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4422 = \"mhlo.dot_general\"(%4421, %4420) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4423 = \"mhlo.slice\"(%4351) {limit_indices = dense<192> : tensor<1xi64>, start_indices = dense<144> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4424 = \"mhlo.slice\"(%4355) {limit_indices = dense<[1, 192]> : tensor<2xi64>, start_indices = dense<[0, 144]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4425 = \"mhlo.dot_general\"(%4424, %4423) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4426 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4427 = mhlo.sqrt %4426 : tensor<f32>\n",
      "    %4428 = mhlo.convert %4427 : tensor<f32>\n",
      "    %4429 = \"mhlo.broadcast_in_dim\"(%4428) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4430 = mhlo.divide %4425, %4429 : tensor<1xf32>\n",
      "    %4431 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4432 = mhlo.reduce(%4430 init: %4431) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4433 = \"mhlo.broadcast_in_dim\"(%4432) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4434 = mhlo.subtract %4430, %4433 : tensor<1xf32>\n",
      "    %4435 = mhlo.exponential %4434 : tensor<1xf32>\n",
      "    %4436 = \"mhlo.slice\"(%4435) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4437 = mhlo.reshape %4436 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4438 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4439 = mhlo.add %4438, %4437 : tensor<f32>\n",
      "    %4440 = \"mhlo.broadcast_in_dim\"(%4439) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4441 = mhlo.divide %4435, %4440 : tensor<1xf32>\n",
      "    %4442 = \"mhlo.slice\"(%4359) {limit_indices = dense<[1, 192]> : tensor<2xi64>, start_indices = dense<[0, 144]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4443 = \"mhlo.dot_general\"(%4442, %4441) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4444 = \"mhlo.slice\"(%4351) {limit_indices = dense<240> : tensor<1xi64>, start_indices = dense<192> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4445 = \"mhlo.slice\"(%4355) {limit_indices = dense<[1, 240]> : tensor<2xi64>, start_indices = dense<[0, 192]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4446 = \"mhlo.dot_general\"(%4445, %4444) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4447 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4448 = mhlo.sqrt %4447 : tensor<f32>\n",
      "    %4449 = mhlo.convert %4448 : tensor<f32>\n",
      "    %4450 = \"mhlo.broadcast_in_dim\"(%4449) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4451 = mhlo.divide %4446, %4450 : tensor<1xf32>\n",
      "    %4452 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4453 = mhlo.reduce(%4451 init: %4452) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4454 = \"mhlo.broadcast_in_dim\"(%4453) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4455 = mhlo.subtract %4451, %4454 : tensor<1xf32>\n",
      "    %4456 = mhlo.exponential %4455 : tensor<1xf32>\n",
      "    %4457 = \"mhlo.slice\"(%4456) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4458 = mhlo.reshape %4457 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4459 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4460 = mhlo.add %4459, %4458 : tensor<f32>\n",
      "    %4461 = \"mhlo.broadcast_in_dim\"(%4460) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4462 = mhlo.divide %4456, %4461 : tensor<1xf32>\n",
      "    %4463 = \"mhlo.slice\"(%4359) {limit_indices = dense<[1, 240]> : tensor<2xi64>, start_indices = dense<[0, 192]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4464 = \"mhlo.dot_general\"(%4463, %4462) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4465 = \"mhlo.slice\"(%4351) {limit_indices = dense<288> : tensor<1xi64>, start_indices = dense<240> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4466 = \"mhlo.slice\"(%4355) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<[0, 240]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4467 = \"mhlo.dot_general\"(%4466, %4465) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4468 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4469 = mhlo.sqrt %4468 : tensor<f32>\n",
      "    %4470 = mhlo.convert %4469 : tensor<f32>\n",
      "    %4471 = \"mhlo.broadcast_in_dim\"(%4470) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4472 = mhlo.divide %4467, %4471 : tensor<1xf32>\n",
      "    %4473 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4474 = mhlo.reduce(%4472 init: %4473) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4475 = \"mhlo.broadcast_in_dim\"(%4474) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4476 = mhlo.subtract %4472, %4475 : tensor<1xf32>\n",
      "    %4477 = mhlo.exponential %4476 : tensor<1xf32>\n",
      "    %4478 = \"mhlo.slice\"(%4477) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4479 = mhlo.reshape %4478 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4480 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4481 = mhlo.add %4480, %4479 : tensor<f32>\n",
      "    %4482 = \"mhlo.broadcast_in_dim\"(%4481) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4483 = mhlo.divide %4477, %4482 : tensor<1xf32>\n",
      "    %4484 = \"mhlo.slice\"(%4359) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<[0, 240]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4485 = \"mhlo.dot_general\"(%4484, %4483) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4486 = \"mhlo.concatenate\"(%4380, %4401, %4422, %4443, %4464, %4485) {dimension = 0 : i64} : (tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>) -> tensor<288xf32>\n",
      "    %4487 = \"mhlo.slice\"(%arg10) {limit_indices = dense<[3, 288, 288]> : tensor<3xi64>, start_indices = dense<[2, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4488 = mhlo.reshape %4487 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4489 = \"mhlo.dot_general\"(%4488, %4486) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4490 = mhlo.add %4323, %4489 : tensor<288xf32>\n",
      "    %4491 = \"mhlo.slice\"(%arg2) {limit_indices = dense<[3, 288]> : tensor<2xi64>, start_indices = dense<[2, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<6x288xf32>) -> tensor<1x288xf32>\n",
      "    %4492 = mhlo.reshape %4491 : (tensor<1x288xf32>) -> tensor<288xf32>\n",
      "    %4493 = \"mhlo.dot_general\"(%4490, %4490) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288xf32>, tensor<288xf32>) -> tensor<f32>\n",
      "    %4494 = mhlo.constant dense<2.880000e+02> : tensor<f32>\n",
      "    %4495 = mhlo.divide %4493, %4494 : tensor<f32>\n",
      "    %4496 = mhlo.constant dense<9.99999974E-6> : tensor<f32>\n",
      "    %4497 = mhlo.add %4495, %4496 : tensor<f32>\n",
      "    %4498 = mhlo.sqrt %4497 : tensor<f32>\n",
      "    %4499 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4500 = mhlo.divide %4499, %4498 : tensor<f32>\n",
      "    %4501 = mhlo.multiply %4492, %4490 : tensor<288xf32>\n",
      "    %4502 = \"mhlo.broadcast_in_dim\"(%4500) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<288xf32>\n",
      "    %4503 = mhlo.multiply %4501, %4502 : tensor<288xf32>\n",
      "    %4504 = \"mhlo.slice\"(%arg5) {limit_indices = dense<[3, 768, 288]> : tensor<3xi64>, start_indices = dense<[2, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x768x288xf32>) -> tensor<1x768x288xf32>\n",
      "    %4505 = mhlo.reshape %4504 : (tensor<1x768x288xf32>) -> tensor<768x288xf32>\n",
      "    %4506 = \"mhlo.dot_general\"(%4505, %4503) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<768x288xf32>, tensor<288xf32>) -> tensor<768xf32>\n",
      "    %4507 = \"mhlo.slice\"(%arg7) {limit_indices = dense<[3, 768, 288]> : tensor<3xi64>, start_indices = dense<[2, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x768x288xf32>) -> tensor<1x768x288xf32>\n",
      "    %4508 = mhlo.reshape %4507 : (tensor<1x768x288xf32>) -> tensor<768x288xf32>\n",
      "    %4509 = \"mhlo.dot_general\"(%4508, %4503) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<768x288xf32>, tensor<288xf32>) -> tensor<768xf32>\n",
      "    %4510 = mhlo.negate %4506 : tensor<768xf32>\n",
      "    %4511 = mhlo.exponential %4510 : tensor<768xf32>\n",
      "    %4512 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4513 = \"mhlo.broadcast_in_dim\"(%4512) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<768xf32>\n",
      "    %4514 = mhlo.add %4513, %4511 : tensor<768xf32>\n",
      "    %4515 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4516 = \"mhlo.broadcast_in_dim\"(%4515) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<768xf32>\n",
      "    %4517 = mhlo.divide %4516, %4514 : tensor<768xf32>\n",
      "    %4518 = mhlo.multiply %4506, %4517 : tensor<768xf32>\n",
      "    %4519 = mhlo.multiply %4518, %4509 : tensor<768xf32>\n",
      "    %4520 = \"mhlo.slice\"(%arg6) {limit_indices = dense<[3, 288, 768]> : tensor<3xi64>, start_indices = dense<[2, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x768xf32>) -> tensor<1x288x768xf32>\n",
      "    %4521 = mhlo.reshape %4520 : (tensor<1x288x768xf32>) -> tensor<288x768xf32>\n",
      "    %4522 = \"mhlo.dot_general\"(%4521, %4519) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x768xf32>, tensor<768xf32>) -> tensor<288xf32>\n",
      "    %4523 = mhlo.add %4490, %4522 : tensor<288xf32>\n",
      "    %4524 = \"mhlo.slice\"(%arg1) {limit_indices = dense<[4, 288]> : tensor<2xi64>, start_indices = dense<[3, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<6x288xf32>) -> tensor<1x288xf32>\n",
      "    %4525 = mhlo.reshape %4524 : (tensor<1x288xf32>) -> tensor<288xf32>\n",
      "    %4526 = \"mhlo.dot_general\"(%4523, %4523) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288xf32>, tensor<288xf32>) -> tensor<f32>\n",
      "    %4527 = mhlo.constant dense<2.880000e+02> : tensor<f32>\n",
      "    %4528 = mhlo.divide %4526, %4527 : tensor<f32>\n",
      "    %4529 = mhlo.constant dense<9.99999974E-6> : tensor<f32>\n",
      "    %4530 = mhlo.add %4528, %4529 : tensor<f32>\n",
      "    %4531 = mhlo.sqrt %4530 : tensor<f32>\n",
      "    %4532 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4533 = mhlo.divide %4532, %4531 : tensor<f32>\n",
      "    %4534 = mhlo.multiply %4525, %4523 : tensor<288xf32>\n",
      "    %4535 = \"mhlo.broadcast_in_dim\"(%4533) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<288xf32>\n",
      "    %4536 = mhlo.multiply %4534, %4535 : tensor<288xf32>\n",
      "    %4537 = \"mhlo.slice\"(%arg11) {limit_indices = dense<[4, 288, 288]> : tensor<3xi64>, start_indices = dense<[3, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4538 = mhlo.reshape %4537 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4539 = \"mhlo.dot_general\"(%4538, %4536) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4540 = \"mhlo.slice\"(%arg9) {limit_indices = dense<[4, 288, 288]> : tensor<3xi64>, start_indices = dense<[3, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4541 = mhlo.reshape %4540 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4542 = \"mhlo.dot_general\"(%4541, %4536) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4543 = \"mhlo.slice\"(%arg12) {limit_indices = dense<[4, 288, 288]> : tensor<3xi64>, start_indices = dense<[3, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4544 = mhlo.reshape %4543 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4545 = \"mhlo.dot_general\"(%4544, %4536) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4546 = mhlo.reshape %4539 : (tensor<288xf32>) -> tensor<144x2xf32>\n",
      "    %4547 = mhlo.reshape %4542 : (tensor<288xf32>) -> tensor<144x2xf32>\n",
      "    %4548 = \"mhlo.dot_general\"(%3923, %4547) {dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>\n",
      "    %4549 = mhlo.reshape %4548 : (tensor<144x2xf32>) -> tensor<288xf32>\n",
      "    %4550 = \"mhlo.dot_general\"(%3769, %4546) {dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>\n",
      "    %4551 = mhlo.reshape %4550 : (tensor<144x2xf32>) -> tensor<288xf32>\n",
      "    %4552 = \"mhlo.slice\"(%arg13) {limit_indices = dense<[4, 0, 288]> : tensor<3xi64>, start_indices = dense<[3, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x0x288xf32>) -> tensor<1x0x288xf32>\n",
      "    %4553 = mhlo.reshape %4552 : (tensor<1x0x288xf32>) -> tensor<0x288xf32>\n",
      "    %4554 = mhlo.reshape %4549 : (tensor<288xf32>) -> tensor<1x288xf32>\n",
      "    %4555 = call @append(%4553, %4554) : (tensor<0x288xf32>, tensor<1x288xf32>) -> tensor<1x288xf32>\n",
      "    %4556 = \"mhlo.slice\"(%arg14) {limit_indices = dense<[4, 0, 288]> : tensor<3xi64>, start_indices = dense<[3, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x0x288xf32>) -> tensor<1x0x288xf32>\n",
      "    %4557 = mhlo.reshape %4556 : (tensor<1x0x288xf32>) -> tensor<0x288xf32>\n",
      "    %4558 = mhlo.reshape %4545 : (tensor<288xf32>) -> tensor<1x288xf32>\n",
      "    %4559 = call @append(%4557, %4558) : (tensor<0x288xf32>, tensor<1x288xf32>) -> tensor<1x288xf32>\n",
      "    %4560 = \"mhlo.slice\"(%4551) {limit_indices = dense<48> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4561 = \"mhlo.slice\"(%4555) {limit_indices = dense<[1, 48]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4562 = \"mhlo.dot_general\"(%4561, %4560) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4563 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4564 = mhlo.sqrt %4563 : tensor<f32>\n",
      "    %4565 = mhlo.convert %4564 : tensor<f32>\n",
      "    %4566 = \"mhlo.broadcast_in_dim\"(%4565) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4567 = mhlo.divide %4562, %4566 : tensor<1xf32>\n",
      "    %4568 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4569 = mhlo.reduce(%4567 init: %4568) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4570 = \"mhlo.broadcast_in_dim\"(%4569) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4571 = mhlo.subtract %4567, %4570 : tensor<1xf32>\n",
      "    %4572 = mhlo.exponential %4571 : tensor<1xf32>\n",
      "    %4573 = \"mhlo.slice\"(%4572) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4574 = mhlo.reshape %4573 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4575 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4576 = mhlo.add %4575, %4574 : tensor<f32>\n",
      "    %4577 = \"mhlo.broadcast_in_dim\"(%4576) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4578 = mhlo.divide %4572, %4577 : tensor<1xf32>\n",
      "    %4579 = \"mhlo.slice\"(%4559) {limit_indices = dense<[1, 48]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4580 = \"mhlo.dot_general\"(%4579, %4578) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4581 = \"mhlo.slice\"(%4551) {limit_indices = dense<96> : tensor<1xi64>, start_indices = dense<48> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4582 = \"mhlo.slice\"(%4555) {limit_indices = dense<[1, 96]> : tensor<2xi64>, start_indices = dense<[0, 48]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4583 = \"mhlo.dot_general\"(%4582, %4581) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4584 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4585 = mhlo.sqrt %4584 : tensor<f32>\n",
      "    %4586 = mhlo.convert %4585 : tensor<f32>\n",
      "    %4587 = \"mhlo.broadcast_in_dim\"(%4586) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4588 = mhlo.divide %4583, %4587 : tensor<1xf32>\n",
      "    %4589 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4590 = mhlo.reduce(%4588 init: %4589) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4591 = \"mhlo.broadcast_in_dim\"(%4590) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4592 = mhlo.subtract %4588, %4591 : tensor<1xf32>\n",
      "    %4593 = mhlo.exponential %4592 : tensor<1xf32>\n",
      "    %4594 = \"mhlo.slice\"(%4593) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4595 = mhlo.reshape %4594 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4596 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4597 = mhlo.add %4596, %4595 : tensor<f32>\n",
      "    %4598 = \"mhlo.broadcast_in_dim\"(%4597) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4599 = mhlo.divide %4593, %4598 : tensor<1xf32>\n",
      "    %4600 = \"mhlo.slice\"(%4559) {limit_indices = dense<[1, 96]> : tensor<2xi64>, start_indices = dense<[0, 48]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4601 = \"mhlo.dot_general\"(%4600, %4599) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4602 = \"mhlo.slice\"(%4551) {limit_indices = dense<144> : tensor<1xi64>, start_indices = dense<96> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4603 = \"mhlo.slice\"(%4555) {limit_indices = dense<[1, 144]> : tensor<2xi64>, start_indices = dense<[0, 96]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4604 = \"mhlo.dot_general\"(%4603, %4602) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4605 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4606 = mhlo.sqrt %4605 : tensor<f32>\n",
      "    %4607 = mhlo.convert %4606 : tensor<f32>\n",
      "    %4608 = \"mhlo.broadcast_in_dim\"(%4607) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4609 = mhlo.divide %4604, %4608 : tensor<1xf32>\n",
      "    %4610 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4611 = mhlo.reduce(%4609 init: %4610) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4612 = \"mhlo.broadcast_in_dim\"(%4611) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4613 = mhlo.subtract %4609, %4612 : tensor<1xf32>\n",
      "    %4614 = mhlo.exponential %4613 : tensor<1xf32>\n",
      "    %4615 = \"mhlo.slice\"(%4614) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4616 = mhlo.reshape %4615 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4617 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4618 = mhlo.add %4617, %4616 : tensor<f32>\n",
      "    %4619 = \"mhlo.broadcast_in_dim\"(%4618) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4620 = mhlo.divide %4614, %4619 : tensor<1xf32>\n",
      "    %4621 = \"mhlo.slice\"(%4559) {limit_indices = dense<[1, 144]> : tensor<2xi64>, start_indices = dense<[0, 96]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4622 = \"mhlo.dot_general\"(%4621, %4620) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4623 = \"mhlo.slice\"(%4551) {limit_indices = dense<192> : tensor<1xi64>, start_indices = dense<144> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4624 = \"mhlo.slice\"(%4555) {limit_indices = dense<[1, 192]> : tensor<2xi64>, start_indices = dense<[0, 144]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4625 = \"mhlo.dot_general\"(%4624, %4623) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4626 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4627 = mhlo.sqrt %4626 : tensor<f32>\n",
      "    %4628 = mhlo.convert %4627 : tensor<f32>\n",
      "    %4629 = \"mhlo.broadcast_in_dim\"(%4628) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4630 = mhlo.divide %4625, %4629 : tensor<1xf32>\n",
      "    %4631 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4632 = mhlo.reduce(%4630 init: %4631) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4633 = \"mhlo.broadcast_in_dim\"(%4632) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4634 = mhlo.subtract %4630, %4633 : tensor<1xf32>\n",
      "    %4635 = mhlo.exponential %4634 : tensor<1xf32>\n",
      "    %4636 = \"mhlo.slice\"(%4635) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4637 = mhlo.reshape %4636 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4638 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4639 = mhlo.add %4638, %4637 : tensor<f32>\n",
      "    %4640 = \"mhlo.broadcast_in_dim\"(%4639) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4641 = mhlo.divide %4635, %4640 : tensor<1xf32>\n",
      "    %4642 = \"mhlo.slice\"(%4559) {limit_indices = dense<[1, 192]> : tensor<2xi64>, start_indices = dense<[0, 144]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4643 = \"mhlo.dot_general\"(%4642, %4641) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4644 = \"mhlo.slice\"(%4551) {limit_indices = dense<240> : tensor<1xi64>, start_indices = dense<192> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4645 = \"mhlo.slice\"(%4555) {limit_indices = dense<[1, 240]> : tensor<2xi64>, start_indices = dense<[0, 192]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4646 = \"mhlo.dot_general\"(%4645, %4644) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4647 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4648 = mhlo.sqrt %4647 : tensor<f32>\n",
      "    %4649 = mhlo.convert %4648 : tensor<f32>\n",
      "    %4650 = \"mhlo.broadcast_in_dim\"(%4649) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4651 = mhlo.divide %4646, %4650 : tensor<1xf32>\n",
      "    %4652 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4653 = mhlo.reduce(%4651 init: %4652) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4654 = \"mhlo.broadcast_in_dim\"(%4653) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4655 = mhlo.subtract %4651, %4654 : tensor<1xf32>\n",
      "    %4656 = mhlo.exponential %4655 : tensor<1xf32>\n",
      "    %4657 = \"mhlo.slice\"(%4656) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4658 = mhlo.reshape %4657 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4659 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4660 = mhlo.add %4659, %4658 : tensor<f32>\n",
      "    %4661 = \"mhlo.broadcast_in_dim\"(%4660) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4662 = mhlo.divide %4656, %4661 : tensor<1xf32>\n",
      "    %4663 = \"mhlo.slice\"(%4559) {limit_indices = dense<[1, 240]> : tensor<2xi64>, start_indices = dense<[0, 192]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4664 = \"mhlo.dot_general\"(%4663, %4662) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4665 = \"mhlo.slice\"(%4551) {limit_indices = dense<288> : tensor<1xi64>, start_indices = dense<240> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4666 = \"mhlo.slice\"(%4555) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<[0, 240]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4667 = \"mhlo.dot_general\"(%4666, %4665) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4668 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4669 = mhlo.sqrt %4668 : tensor<f32>\n",
      "    %4670 = mhlo.convert %4669 : tensor<f32>\n",
      "    %4671 = \"mhlo.broadcast_in_dim\"(%4670) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4672 = mhlo.divide %4667, %4671 : tensor<1xf32>\n",
      "    %4673 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4674 = mhlo.reduce(%4672 init: %4673) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4675 = \"mhlo.broadcast_in_dim\"(%4674) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4676 = mhlo.subtract %4672, %4675 : tensor<1xf32>\n",
      "    %4677 = mhlo.exponential %4676 : tensor<1xf32>\n",
      "    %4678 = \"mhlo.slice\"(%4677) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4679 = mhlo.reshape %4678 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4680 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4681 = mhlo.add %4680, %4679 : tensor<f32>\n",
      "    %4682 = \"mhlo.broadcast_in_dim\"(%4681) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4683 = mhlo.divide %4677, %4682 : tensor<1xf32>\n",
      "    %4684 = \"mhlo.slice\"(%4559) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<[0, 240]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4685 = \"mhlo.dot_general\"(%4684, %4683) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4686 = \"mhlo.concatenate\"(%4580, %4601, %4622, %4643, %4664, %4685) {dimension = 0 : i64} : (tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>) -> tensor<288xf32>\n",
      "    %4687 = \"mhlo.slice\"(%arg10) {limit_indices = dense<[4, 288, 288]> : tensor<3xi64>, start_indices = dense<[3, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4688 = mhlo.reshape %4687 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4689 = \"mhlo.dot_general\"(%4688, %4686) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4690 = mhlo.add %4523, %4689 : tensor<288xf32>\n",
      "    %4691 = \"mhlo.slice\"(%arg2) {limit_indices = dense<[4, 288]> : tensor<2xi64>, start_indices = dense<[3, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<6x288xf32>) -> tensor<1x288xf32>\n",
      "    %4692 = mhlo.reshape %4691 : (tensor<1x288xf32>) -> tensor<288xf32>\n",
      "    %4693 = \"mhlo.dot_general\"(%4690, %4690) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288xf32>, tensor<288xf32>) -> tensor<f32>\n",
      "    %4694 = mhlo.constant dense<2.880000e+02> : tensor<f32>\n",
      "    %4695 = mhlo.divide %4693, %4694 : tensor<f32>\n",
      "    %4696 = mhlo.constant dense<9.99999974E-6> : tensor<f32>\n",
      "    %4697 = mhlo.add %4695, %4696 : tensor<f32>\n",
      "    %4698 = mhlo.sqrt %4697 : tensor<f32>\n",
      "    %4699 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4700 = mhlo.divide %4699, %4698 : tensor<f32>\n",
      "    %4701 = mhlo.multiply %4692, %4690 : tensor<288xf32>\n",
      "    %4702 = \"mhlo.broadcast_in_dim\"(%4700) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<288xf32>\n",
      "    %4703 = mhlo.multiply %4701, %4702 : tensor<288xf32>\n",
      "    %4704 = \"mhlo.slice\"(%arg5) {limit_indices = dense<[4, 768, 288]> : tensor<3xi64>, start_indices = dense<[3, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x768x288xf32>) -> tensor<1x768x288xf32>\n",
      "    %4705 = mhlo.reshape %4704 : (tensor<1x768x288xf32>) -> tensor<768x288xf32>\n",
      "    %4706 = \"mhlo.dot_general\"(%4705, %4703) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<768x288xf32>, tensor<288xf32>) -> tensor<768xf32>\n",
      "    %4707 = \"mhlo.slice\"(%arg7) {limit_indices = dense<[4, 768, 288]> : tensor<3xi64>, start_indices = dense<[3, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x768x288xf32>) -> tensor<1x768x288xf32>\n",
      "    %4708 = mhlo.reshape %4707 : (tensor<1x768x288xf32>) -> tensor<768x288xf32>\n",
      "    %4709 = \"mhlo.dot_general\"(%4708, %4703) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<768x288xf32>, tensor<288xf32>) -> tensor<768xf32>\n",
      "    %4710 = mhlo.negate %4706 : tensor<768xf32>\n",
      "    %4711 = mhlo.exponential %4710 : tensor<768xf32>\n",
      "    %4712 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4713 = \"mhlo.broadcast_in_dim\"(%4712) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<768xf32>\n",
      "    %4714 = mhlo.add %4713, %4711 : tensor<768xf32>\n",
      "    %4715 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4716 = \"mhlo.broadcast_in_dim\"(%4715) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<768xf32>\n",
      "    %4717 = mhlo.divide %4716, %4714 : tensor<768xf32>\n",
      "    %4718 = mhlo.multiply %4706, %4717 : tensor<768xf32>\n",
      "    %4719 = mhlo.multiply %4718, %4709 : tensor<768xf32>\n",
      "    %4720 = \"mhlo.slice\"(%arg6) {limit_indices = dense<[4, 288, 768]> : tensor<3xi64>, start_indices = dense<[3, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x768xf32>) -> tensor<1x288x768xf32>\n",
      "    %4721 = mhlo.reshape %4720 : (tensor<1x288x768xf32>) -> tensor<288x768xf32>\n",
      "    %4722 = \"mhlo.dot_general\"(%4721, %4719) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x768xf32>, tensor<768xf32>) -> tensor<288xf32>\n",
      "    %4723 = mhlo.add %4690, %4722 : tensor<288xf32>\n",
      "    %4724 = \"mhlo.slice\"(%arg1) {limit_indices = dense<[5, 288]> : tensor<2xi64>, start_indices = dense<[4, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<6x288xf32>) -> tensor<1x288xf32>\n",
      "    %4725 = mhlo.reshape %4724 : (tensor<1x288xf32>) -> tensor<288xf32>\n",
      "    %4726 = \"mhlo.dot_general\"(%4723, %4723) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288xf32>, tensor<288xf32>) -> tensor<f32>\n",
      "    %4727 = mhlo.constant dense<2.880000e+02> : tensor<f32>\n",
      "    %4728 = mhlo.divide %4726, %4727 : tensor<f32>\n",
      "    %4729 = mhlo.constant dense<9.99999974E-6> : tensor<f32>\n",
      "    %4730 = mhlo.add %4728, %4729 : tensor<f32>\n",
      "    %4731 = mhlo.sqrt %4730 : tensor<f32>\n",
      "    %4732 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4733 = mhlo.divide %4732, %4731 : tensor<f32>\n",
      "    %4734 = mhlo.multiply %4725, %4723 : tensor<288xf32>\n",
      "    %4735 = \"mhlo.broadcast_in_dim\"(%4733) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<288xf32>\n",
      "    %4736 = mhlo.multiply %4734, %4735 : tensor<288xf32>\n",
      "    %4737 = \"mhlo.slice\"(%arg11) {limit_indices = dense<[5, 288, 288]> : tensor<3xi64>, start_indices = dense<[4, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4738 = mhlo.reshape %4737 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4739 = \"mhlo.dot_general\"(%4738, %4736) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4740 = \"mhlo.slice\"(%arg9) {limit_indices = dense<[5, 288, 288]> : tensor<3xi64>, start_indices = dense<[4, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4741 = mhlo.reshape %4740 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4742 = \"mhlo.dot_general\"(%4741, %4736) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4743 = \"mhlo.slice\"(%arg12) {limit_indices = dense<[5, 288, 288]> : tensor<3xi64>, start_indices = dense<[4, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4744 = mhlo.reshape %4743 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4745 = \"mhlo.dot_general\"(%4744, %4736) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4746 = mhlo.reshape %4739 : (tensor<288xf32>) -> tensor<144x2xf32>\n",
      "    %4747 = mhlo.reshape %4742 : (tensor<288xf32>) -> tensor<144x2xf32>\n",
      "    %4748 = \"mhlo.dot_general\"(%3923, %4747) {dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>\n",
      "    %4749 = mhlo.reshape %4748 : (tensor<144x2xf32>) -> tensor<288xf32>\n",
      "    %4750 = \"mhlo.dot_general\"(%3769, %4746) {dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>\n",
      "    %4751 = mhlo.reshape %4750 : (tensor<144x2xf32>) -> tensor<288xf32>\n",
      "    %4752 = \"mhlo.slice\"(%arg13) {limit_indices = dense<[5, 0, 288]> : tensor<3xi64>, start_indices = dense<[4, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x0x288xf32>) -> tensor<1x0x288xf32>\n",
      "    %4753 = mhlo.reshape %4752 : (tensor<1x0x288xf32>) -> tensor<0x288xf32>\n",
      "    %4754 = mhlo.reshape %4749 : (tensor<288xf32>) -> tensor<1x288xf32>\n",
      "    %4755 = call @append(%4753, %4754) : (tensor<0x288xf32>, tensor<1x288xf32>) -> tensor<1x288xf32>\n",
      "    %4756 = \"mhlo.slice\"(%arg14) {limit_indices = dense<[5, 0, 288]> : tensor<3xi64>, start_indices = dense<[4, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x0x288xf32>) -> tensor<1x0x288xf32>\n",
      "    %4757 = mhlo.reshape %4756 : (tensor<1x0x288xf32>) -> tensor<0x288xf32>\n",
      "    %4758 = mhlo.reshape %4745 : (tensor<288xf32>) -> tensor<1x288xf32>\n",
      "    %4759 = call @append(%4757, %4758) : (tensor<0x288xf32>, tensor<1x288xf32>) -> tensor<1x288xf32>\n",
      "    %4760 = \"mhlo.slice\"(%4751) {limit_indices = dense<48> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4761 = \"mhlo.slice\"(%4755) {limit_indices = dense<[1, 48]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4762 = \"mhlo.dot_general\"(%4761, %4760) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4763 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4764 = mhlo.sqrt %4763 : tensor<f32>\n",
      "    %4765 = mhlo.convert %4764 : tensor<f32>\n",
      "    %4766 = \"mhlo.broadcast_in_dim\"(%4765) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4767 = mhlo.divide %4762, %4766 : tensor<1xf32>\n",
      "    %4768 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4769 = mhlo.reduce(%4767 init: %4768) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4770 = \"mhlo.broadcast_in_dim\"(%4769) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4771 = mhlo.subtract %4767, %4770 : tensor<1xf32>\n",
      "    %4772 = mhlo.exponential %4771 : tensor<1xf32>\n",
      "    %4773 = \"mhlo.slice\"(%4772) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4774 = mhlo.reshape %4773 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4775 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4776 = mhlo.add %4775, %4774 : tensor<f32>\n",
      "    %4777 = \"mhlo.broadcast_in_dim\"(%4776) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4778 = mhlo.divide %4772, %4777 : tensor<1xf32>\n",
      "    %4779 = \"mhlo.slice\"(%4759) {limit_indices = dense<[1, 48]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4780 = \"mhlo.dot_general\"(%4779, %4778) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4781 = \"mhlo.slice\"(%4751) {limit_indices = dense<96> : tensor<1xi64>, start_indices = dense<48> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4782 = \"mhlo.slice\"(%4755) {limit_indices = dense<[1, 96]> : tensor<2xi64>, start_indices = dense<[0, 48]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4783 = \"mhlo.dot_general\"(%4782, %4781) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4784 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4785 = mhlo.sqrt %4784 : tensor<f32>\n",
      "    %4786 = mhlo.convert %4785 : tensor<f32>\n",
      "    %4787 = \"mhlo.broadcast_in_dim\"(%4786) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4788 = mhlo.divide %4783, %4787 : tensor<1xf32>\n",
      "    %4789 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4790 = mhlo.reduce(%4788 init: %4789) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4791 = \"mhlo.broadcast_in_dim\"(%4790) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4792 = mhlo.subtract %4788, %4791 : tensor<1xf32>\n",
      "    %4793 = mhlo.exponential %4792 : tensor<1xf32>\n",
      "    %4794 = \"mhlo.slice\"(%4793) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4795 = mhlo.reshape %4794 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4796 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4797 = mhlo.add %4796, %4795 : tensor<f32>\n",
      "    %4798 = \"mhlo.broadcast_in_dim\"(%4797) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4799 = mhlo.divide %4793, %4798 : tensor<1xf32>\n",
      "    %4800 = \"mhlo.slice\"(%4759) {limit_indices = dense<[1, 96]> : tensor<2xi64>, start_indices = dense<[0, 48]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4801 = \"mhlo.dot_general\"(%4800, %4799) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4802 = \"mhlo.slice\"(%4751) {limit_indices = dense<144> : tensor<1xi64>, start_indices = dense<96> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4803 = \"mhlo.slice\"(%4755) {limit_indices = dense<[1, 144]> : tensor<2xi64>, start_indices = dense<[0, 96]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4804 = \"mhlo.dot_general\"(%4803, %4802) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4805 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4806 = mhlo.sqrt %4805 : tensor<f32>\n",
      "    %4807 = mhlo.convert %4806 : tensor<f32>\n",
      "    %4808 = \"mhlo.broadcast_in_dim\"(%4807) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4809 = mhlo.divide %4804, %4808 : tensor<1xf32>\n",
      "    %4810 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4811 = mhlo.reduce(%4809 init: %4810) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4812 = \"mhlo.broadcast_in_dim\"(%4811) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4813 = mhlo.subtract %4809, %4812 : tensor<1xf32>\n",
      "    %4814 = mhlo.exponential %4813 : tensor<1xf32>\n",
      "    %4815 = \"mhlo.slice\"(%4814) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4816 = mhlo.reshape %4815 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4817 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4818 = mhlo.add %4817, %4816 : tensor<f32>\n",
      "    %4819 = \"mhlo.broadcast_in_dim\"(%4818) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4820 = mhlo.divide %4814, %4819 : tensor<1xf32>\n",
      "    %4821 = \"mhlo.slice\"(%4759) {limit_indices = dense<[1, 144]> : tensor<2xi64>, start_indices = dense<[0, 96]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4822 = \"mhlo.dot_general\"(%4821, %4820) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4823 = \"mhlo.slice\"(%4751) {limit_indices = dense<192> : tensor<1xi64>, start_indices = dense<144> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4824 = \"mhlo.slice\"(%4755) {limit_indices = dense<[1, 192]> : tensor<2xi64>, start_indices = dense<[0, 144]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4825 = \"mhlo.dot_general\"(%4824, %4823) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4826 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4827 = mhlo.sqrt %4826 : tensor<f32>\n",
      "    %4828 = mhlo.convert %4827 : tensor<f32>\n",
      "    %4829 = \"mhlo.broadcast_in_dim\"(%4828) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4830 = mhlo.divide %4825, %4829 : tensor<1xf32>\n",
      "    %4831 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4832 = mhlo.reduce(%4830 init: %4831) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4833 = \"mhlo.broadcast_in_dim\"(%4832) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4834 = mhlo.subtract %4830, %4833 : tensor<1xf32>\n",
      "    %4835 = mhlo.exponential %4834 : tensor<1xf32>\n",
      "    %4836 = \"mhlo.slice\"(%4835) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4837 = mhlo.reshape %4836 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4838 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4839 = mhlo.add %4838, %4837 : tensor<f32>\n",
      "    %4840 = \"mhlo.broadcast_in_dim\"(%4839) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4841 = mhlo.divide %4835, %4840 : tensor<1xf32>\n",
      "    %4842 = \"mhlo.slice\"(%4759) {limit_indices = dense<[1, 192]> : tensor<2xi64>, start_indices = dense<[0, 144]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4843 = \"mhlo.dot_general\"(%4842, %4841) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4844 = \"mhlo.slice\"(%4751) {limit_indices = dense<240> : tensor<1xi64>, start_indices = dense<192> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4845 = \"mhlo.slice\"(%4755) {limit_indices = dense<[1, 240]> : tensor<2xi64>, start_indices = dense<[0, 192]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4846 = \"mhlo.dot_general\"(%4845, %4844) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4847 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4848 = mhlo.sqrt %4847 : tensor<f32>\n",
      "    %4849 = mhlo.convert %4848 : tensor<f32>\n",
      "    %4850 = \"mhlo.broadcast_in_dim\"(%4849) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4851 = mhlo.divide %4846, %4850 : tensor<1xf32>\n",
      "    %4852 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4853 = mhlo.reduce(%4851 init: %4852) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4854 = \"mhlo.broadcast_in_dim\"(%4853) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4855 = mhlo.subtract %4851, %4854 : tensor<1xf32>\n",
      "    %4856 = mhlo.exponential %4855 : tensor<1xf32>\n",
      "    %4857 = \"mhlo.slice\"(%4856) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4858 = mhlo.reshape %4857 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4859 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4860 = mhlo.add %4859, %4858 : tensor<f32>\n",
      "    %4861 = \"mhlo.broadcast_in_dim\"(%4860) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4862 = mhlo.divide %4856, %4861 : tensor<1xf32>\n",
      "    %4863 = \"mhlo.slice\"(%4759) {limit_indices = dense<[1, 240]> : tensor<2xi64>, start_indices = dense<[0, 192]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4864 = \"mhlo.dot_general\"(%4863, %4862) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4865 = \"mhlo.slice\"(%4751) {limit_indices = dense<288> : tensor<1xi64>, start_indices = dense<240> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4866 = \"mhlo.slice\"(%4755) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<[0, 240]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4867 = \"mhlo.dot_general\"(%4866, %4865) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4868 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4869 = mhlo.sqrt %4868 : tensor<f32>\n",
      "    %4870 = mhlo.convert %4869 : tensor<f32>\n",
      "    %4871 = \"mhlo.broadcast_in_dim\"(%4870) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4872 = mhlo.divide %4867, %4871 : tensor<1xf32>\n",
      "    %4873 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4874 = mhlo.reduce(%4872 init: %4873) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4875 = \"mhlo.broadcast_in_dim\"(%4874) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4876 = mhlo.subtract %4872, %4875 : tensor<1xf32>\n",
      "    %4877 = mhlo.exponential %4876 : tensor<1xf32>\n",
      "    %4878 = \"mhlo.slice\"(%4877) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4879 = mhlo.reshape %4878 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4880 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4881 = mhlo.add %4880, %4879 : tensor<f32>\n",
      "    %4882 = \"mhlo.broadcast_in_dim\"(%4881) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4883 = mhlo.divide %4877, %4882 : tensor<1xf32>\n",
      "    %4884 = \"mhlo.slice\"(%4759) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<[0, 240]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4885 = \"mhlo.dot_general\"(%4884, %4883) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4886 = \"mhlo.concatenate\"(%4780, %4801, %4822, %4843, %4864, %4885) {dimension = 0 : i64} : (tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>) -> tensor<288xf32>\n",
      "    %4887 = \"mhlo.slice\"(%arg10) {limit_indices = dense<[5, 288, 288]> : tensor<3xi64>, start_indices = dense<[4, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4888 = mhlo.reshape %4887 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4889 = \"mhlo.dot_general\"(%4888, %4886) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4890 = mhlo.add %4723, %4889 : tensor<288xf32>\n",
      "    %4891 = \"mhlo.slice\"(%arg2) {limit_indices = dense<[5, 288]> : tensor<2xi64>, start_indices = dense<[4, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<6x288xf32>) -> tensor<1x288xf32>\n",
      "    %4892 = mhlo.reshape %4891 : (tensor<1x288xf32>) -> tensor<288xf32>\n",
      "    %4893 = \"mhlo.dot_general\"(%4890, %4890) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288xf32>, tensor<288xf32>) -> tensor<f32>\n",
      "    %4894 = mhlo.constant dense<2.880000e+02> : tensor<f32>\n",
      "    %4895 = mhlo.divide %4893, %4894 : tensor<f32>\n",
      "    %4896 = mhlo.constant dense<9.99999974E-6> : tensor<f32>\n",
      "    %4897 = mhlo.add %4895, %4896 : tensor<f32>\n",
      "    %4898 = mhlo.sqrt %4897 : tensor<f32>\n",
      "    %4899 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4900 = mhlo.divide %4899, %4898 : tensor<f32>\n",
      "    %4901 = mhlo.multiply %4892, %4890 : tensor<288xf32>\n",
      "    %4902 = \"mhlo.broadcast_in_dim\"(%4900) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<288xf32>\n",
      "    %4903 = mhlo.multiply %4901, %4902 : tensor<288xf32>\n",
      "    %4904 = \"mhlo.slice\"(%arg5) {limit_indices = dense<[5, 768, 288]> : tensor<3xi64>, start_indices = dense<[4, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x768x288xf32>) -> tensor<1x768x288xf32>\n",
      "    %4905 = mhlo.reshape %4904 : (tensor<1x768x288xf32>) -> tensor<768x288xf32>\n",
      "    %4906 = \"mhlo.dot_general\"(%4905, %4903) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<768x288xf32>, tensor<288xf32>) -> tensor<768xf32>\n",
      "    %4907 = \"mhlo.slice\"(%arg7) {limit_indices = dense<[5, 768, 288]> : tensor<3xi64>, start_indices = dense<[4, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x768x288xf32>) -> tensor<1x768x288xf32>\n",
      "    %4908 = mhlo.reshape %4907 : (tensor<1x768x288xf32>) -> tensor<768x288xf32>\n",
      "    %4909 = \"mhlo.dot_general\"(%4908, %4903) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<768x288xf32>, tensor<288xf32>) -> tensor<768xf32>\n",
      "    %4910 = mhlo.negate %4906 : tensor<768xf32>\n",
      "    %4911 = mhlo.exponential %4910 : tensor<768xf32>\n",
      "    %4912 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4913 = \"mhlo.broadcast_in_dim\"(%4912) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<768xf32>\n",
      "    %4914 = mhlo.add %4913, %4911 : tensor<768xf32>\n",
      "    %4915 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4916 = \"mhlo.broadcast_in_dim\"(%4915) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<768xf32>\n",
      "    %4917 = mhlo.divide %4916, %4914 : tensor<768xf32>\n",
      "    %4918 = mhlo.multiply %4906, %4917 : tensor<768xf32>\n",
      "    %4919 = mhlo.multiply %4918, %4909 : tensor<768xf32>\n",
      "    %4920 = \"mhlo.slice\"(%arg6) {limit_indices = dense<[5, 288, 768]> : tensor<3xi64>, start_indices = dense<[4, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x768xf32>) -> tensor<1x288x768xf32>\n",
      "    %4921 = mhlo.reshape %4920 : (tensor<1x288x768xf32>) -> tensor<288x768xf32>\n",
      "    %4922 = \"mhlo.dot_general\"(%4921, %4919) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x768xf32>, tensor<768xf32>) -> tensor<288xf32>\n",
      "    %4923 = mhlo.add %4890, %4922 : tensor<288xf32>\n",
      "    %4924 = \"mhlo.slice\"(%arg1) {limit_indices = dense<[6, 288]> : tensor<2xi64>, start_indices = dense<[5, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<6x288xf32>) -> tensor<1x288xf32>\n",
      "    %4925 = mhlo.reshape %4924 : (tensor<1x288xf32>) -> tensor<288xf32>\n",
      "    %4926 = \"mhlo.dot_general\"(%4923, %4923) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288xf32>, tensor<288xf32>) -> tensor<f32>\n",
      "    %4927 = mhlo.constant dense<2.880000e+02> : tensor<f32>\n",
      "    %4928 = mhlo.divide %4926, %4927 : tensor<f32>\n",
      "    %4929 = mhlo.constant dense<9.99999974E-6> : tensor<f32>\n",
      "    %4930 = mhlo.add %4928, %4929 : tensor<f32>\n",
      "    %4931 = mhlo.sqrt %4930 : tensor<f32>\n",
      "    %4932 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %4933 = mhlo.divide %4932, %4931 : tensor<f32>\n",
      "    %4934 = mhlo.multiply %4925, %4923 : tensor<288xf32>\n",
      "    %4935 = \"mhlo.broadcast_in_dim\"(%4933) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<288xf32>\n",
      "    %4936 = mhlo.multiply %4934, %4935 : tensor<288xf32>\n",
      "    %4937 = \"mhlo.slice\"(%arg11) {limit_indices = dense<[6, 288, 288]> : tensor<3xi64>, start_indices = dense<[5, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4938 = mhlo.reshape %4937 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4939 = \"mhlo.dot_general\"(%4938, %4936) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4940 = \"mhlo.slice\"(%arg9) {limit_indices = dense<[6, 288, 288]> : tensor<3xi64>, start_indices = dense<[5, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4941 = mhlo.reshape %4940 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4942 = \"mhlo.dot_general\"(%4941, %4936) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4943 = \"mhlo.slice\"(%arg12) {limit_indices = dense<[6, 288, 288]> : tensor<3xi64>, start_indices = dense<[5, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %4944 = mhlo.reshape %4943 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %4945 = \"mhlo.dot_general\"(%4944, %4936) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %4946 = mhlo.reshape %4939 : (tensor<288xf32>) -> tensor<144x2xf32>\n",
      "    %4947 = mhlo.reshape %4942 : (tensor<288xf32>) -> tensor<144x2xf32>\n",
      "    %4948 = \"mhlo.dot_general\"(%3923, %4947) {dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>\n",
      "    %4949 = mhlo.reshape %4948 : (tensor<144x2xf32>) -> tensor<288xf32>\n",
      "    %4950 = \"mhlo.dot_general\"(%3769, %4946) {dot_dimension_numbers = #mhlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<144x2x2xf32>, tensor<144x2xf32>) -> tensor<144x2xf32>\n",
      "    %4951 = mhlo.reshape %4950 : (tensor<144x2xf32>) -> tensor<288xf32>\n",
      "    %4952 = \"mhlo.slice\"(%arg13) {limit_indices = dense<[6, 0, 288]> : tensor<3xi64>, start_indices = dense<[5, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x0x288xf32>) -> tensor<1x0x288xf32>\n",
      "    %4953 = mhlo.reshape %4952 : (tensor<1x0x288xf32>) -> tensor<0x288xf32>\n",
      "    %4954 = mhlo.reshape %4949 : (tensor<288xf32>) -> tensor<1x288xf32>\n",
      "    %4955 = call @append(%4953, %4954) : (tensor<0x288xf32>, tensor<1x288xf32>) -> tensor<1x288xf32>\n",
      "    %4956 = \"mhlo.slice\"(%arg14) {limit_indices = dense<[6, 0, 288]> : tensor<3xi64>, start_indices = dense<[5, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x0x288xf32>) -> tensor<1x0x288xf32>\n",
      "    %4957 = mhlo.reshape %4956 : (tensor<1x0x288xf32>) -> tensor<0x288xf32>\n",
      "    %4958 = mhlo.reshape %4945 : (tensor<288xf32>) -> tensor<1x288xf32>\n",
      "    %4959 = call @append(%4957, %4958) : (tensor<0x288xf32>, tensor<1x288xf32>) -> tensor<1x288xf32>\n",
      "    %4960 = \"mhlo.slice\"(%4951) {limit_indices = dense<48> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4961 = \"mhlo.slice\"(%4955) {limit_indices = dense<[1, 48]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4962 = \"mhlo.dot_general\"(%4961, %4960) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4963 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4964 = mhlo.sqrt %4963 : tensor<f32>\n",
      "    %4965 = mhlo.convert %4964 : tensor<f32>\n",
      "    %4966 = \"mhlo.broadcast_in_dim\"(%4965) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4967 = mhlo.divide %4962, %4966 : tensor<1xf32>\n",
      "    %4968 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4969 = mhlo.reduce(%4967 init: %4968) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4970 = \"mhlo.broadcast_in_dim\"(%4969) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4971 = mhlo.subtract %4967, %4970 : tensor<1xf32>\n",
      "    %4972 = mhlo.exponential %4971 : tensor<1xf32>\n",
      "    %4973 = \"mhlo.slice\"(%4972) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4974 = mhlo.reshape %4973 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4975 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4976 = mhlo.add %4975, %4974 : tensor<f32>\n",
      "    %4977 = \"mhlo.broadcast_in_dim\"(%4976) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4978 = mhlo.divide %4972, %4977 : tensor<1xf32>\n",
      "    %4979 = \"mhlo.slice\"(%4959) {limit_indices = dense<[1, 48]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4980 = \"mhlo.dot_general\"(%4979, %4978) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %4981 = \"mhlo.slice\"(%4951) {limit_indices = dense<96> : tensor<1xi64>, start_indices = dense<48> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %4982 = \"mhlo.slice\"(%4955) {limit_indices = dense<[1, 96]> : tensor<2xi64>, start_indices = dense<[0, 48]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %4983 = \"mhlo.dot_general\"(%4982, %4981) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %4984 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %4985 = mhlo.sqrt %4984 : tensor<f32>\n",
      "    %4986 = mhlo.convert %4985 : tensor<f32>\n",
      "    %4987 = \"mhlo.broadcast_in_dim\"(%4986) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4988 = mhlo.divide %4983, %4987 : tensor<1xf32>\n",
      "    %4989 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %4990 = mhlo.reduce(%4988 init: %4989) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %4991 = \"mhlo.broadcast_in_dim\"(%4990) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4992 = mhlo.subtract %4988, %4991 : tensor<1xf32>\n",
      "    %4993 = mhlo.exponential %4992 : tensor<1xf32>\n",
      "    %4994 = \"mhlo.slice\"(%4993) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %4995 = mhlo.reshape %4994 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %4996 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %4997 = mhlo.add %4996, %4995 : tensor<f32>\n",
      "    %4998 = \"mhlo.broadcast_in_dim\"(%4997) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %4999 = mhlo.divide %4993, %4998 : tensor<1xf32>\n",
      "    %5000 = \"mhlo.slice\"(%4959) {limit_indices = dense<[1, 96]> : tensor<2xi64>, start_indices = dense<[0, 48]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %5001 = \"mhlo.dot_general\"(%5000, %4999) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %5002 = \"mhlo.slice\"(%4951) {limit_indices = dense<144> : tensor<1xi64>, start_indices = dense<96> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %5003 = \"mhlo.slice\"(%4955) {limit_indices = dense<[1, 144]> : tensor<2xi64>, start_indices = dense<[0, 96]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %5004 = \"mhlo.dot_general\"(%5003, %5002) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %5005 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %5006 = mhlo.sqrt %5005 : tensor<f32>\n",
      "    %5007 = mhlo.convert %5006 : tensor<f32>\n",
      "    %5008 = \"mhlo.broadcast_in_dim\"(%5007) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %5009 = mhlo.divide %5004, %5008 : tensor<1xf32>\n",
      "    %5010 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %5011 = mhlo.reduce(%5009 init: %5010) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %5012 = \"mhlo.broadcast_in_dim\"(%5011) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %5013 = mhlo.subtract %5009, %5012 : tensor<1xf32>\n",
      "    %5014 = mhlo.exponential %5013 : tensor<1xf32>\n",
      "    %5015 = \"mhlo.slice\"(%5014) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %5016 = mhlo.reshape %5015 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %5017 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %5018 = mhlo.add %5017, %5016 : tensor<f32>\n",
      "    %5019 = \"mhlo.broadcast_in_dim\"(%5018) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %5020 = mhlo.divide %5014, %5019 : tensor<1xf32>\n",
      "    %5021 = \"mhlo.slice\"(%4959) {limit_indices = dense<[1, 144]> : tensor<2xi64>, start_indices = dense<[0, 96]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %5022 = \"mhlo.dot_general\"(%5021, %5020) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %5023 = \"mhlo.slice\"(%4951) {limit_indices = dense<192> : tensor<1xi64>, start_indices = dense<144> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %5024 = \"mhlo.slice\"(%4955) {limit_indices = dense<[1, 192]> : tensor<2xi64>, start_indices = dense<[0, 144]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %5025 = \"mhlo.dot_general\"(%5024, %5023) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %5026 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %5027 = mhlo.sqrt %5026 : tensor<f32>\n",
      "    %5028 = mhlo.convert %5027 : tensor<f32>\n",
      "    %5029 = \"mhlo.broadcast_in_dim\"(%5028) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %5030 = mhlo.divide %5025, %5029 : tensor<1xf32>\n",
      "    %5031 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %5032 = mhlo.reduce(%5030 init: %5031) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %5033 = \"mhlo.broadcast_in_dim\"(%5032) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %5034 = mhlo.subtract %5030, %5033 : tensor<1xf32>\n",
      "    %5035 = mhlo.exponential %5034 : tensor<1xf32>\n",
      "    %5036 = \"mhlo.slice\"(%5035) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %5037 = mhlo.reshape %5036 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %5038 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %5039 = mhlo.add %5038, %5037 : tensor<f32>\n",
      "    %5040 = \"mhlo.broadcast_in_dim\"(%5039) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %5041 = mhlo.divide %5035, %5040 : tensor<1xf32>\n",
      "    %5042 = \"mhlo.slice\"(%4959) {limit_indices = dense<[1, 192]> : tensor<2xi64>, start_indices = dense<[0, 144]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %5043 = \"mhlo.dot_general\"(%5042, %5041) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %5044 = \"mhlo.slice\"(%4951) {limit_indices = dense<240> : tensor<1xi64>, start_indices = dense<192> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %5045 = \"mhlo.slice\"(%4955) {limit_indices = dense<[1, 240]> : tensor<2xi64>, start_indices = dense<[0, 192]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %5046 = \"mhlo.dot_general\"(%5045, %5044) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %5047 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %5048 = mhlo.sqrt %5047 : tensor<f32>\n",
      "    %5049 = mhlo.convert %5048 : tensor<f32>\n",
      "    %5050 = \"mhlo.broadcast_in_dim\"(%5049) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %5051 = mhlo.divide %5046, %5050 : tensor<1xf32>\n",
      "    %5052 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %5053 = mhlo.reduce(%5051 init: %5052) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %5054 = \"mhlo.broadcast_in_dim\"(%5053) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %5055 = mhlo.subtract %5051, %5054 : tensor<1xf32>\n",
      "    %5056 = mhlo.exponential %5055 : tensor<1xf32>\n",
      "    %5057 = \"mhlo.slice\"(%5056) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %5058 = mhlo.reshape %5057 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %5059 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %5060 = mhlo.add %5059, %5058 : tensor<f32>\n",
      "    %5061 = \"mhlo.broadcast_in_dim\"(%5060) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %5062 = mhlo.divide %5056, %5061 : tensor<1xf32>\n",
      "    %5063 = \"mhlo.slice\"(%4959) {limit_indices = dense<[1, 240]> : tensor<2xi64>, start_indices = dense<[0, 192]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %5064 = \"mhlo.dot_general\"(%5063, %5062) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %5065 = \"mhlo.slice\"(%4951) {limit_indices = dense<288> : tensor<1xi64>, start_indices = dense<240> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<288xf32>) -> tensor<48xf32>\n",
      "    %5066 = \"mhlo.slice\"(%4955) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<[0, 240]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %5067 = \"mhlo.dot_general\"(%5066, %5065) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<48xf32>) -> tensor<1xf32>\n",
      "    %5068 = mhlo.constant dense<4.800000e+01> : tensor<f32>\n",
      "    %5069 = mhlo.sqrt %5068 : tensor<f32>\n",
      "    %5070 = mhlo.convert %5069 : tensor<f32>\n",
      "    %5071 = \"mhlo.broadcast_in_dim\"(%5070) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %5072 = mhlo.divide %5067, %5071 : tensor<1xf32>\n",
      "    %5073 = mhlo.constant dense<0xFF800000> : tensor<f32>\n",
      "    %5074 = mhlo.reduce(%5072 init: %5073) applies mhlo.maximum across dimensions = [0] : (tensor<1xf32>, tensor<f32>) -> tensor<f32>\n",
      "    %5075 = \"mhlo.broadcast_in_dim\"(%5074) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %5076 = mhlo.subtract %5072, %5075 : tensor<1xf32>\n",
      "    %5077 = mhlo.exponential %5076 : tensor<1xf32>\n",
      "    %5078 = \"mhlo.slice\"(%5077) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1xf32>) -> tensor<1xf32>\n",
      "    %5079 = mhlo.reshape %5078 : (tensor<1xf32>) -> tensor<f32>\n",
      "    %5080 = mhlo.constant dense<0.000000e+00> : tensor<f32>\n",
      "    %5081 = mhlo.add %5080, %5079 : tensor<f32>\n",
      "    %5082 = \"mhlo.broadcast_in_dim\"(%5081) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1xf32>\n",
      "    %5083 = mhlo.divide %5077, %5082 : tensor<1xf32>\n",
      "    %5084 = \"mhlo.slice\"(%4959) {limit_indices = dense<[1, 288]> : tensor<2xi64>, start_indices = dense<[0, 240]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x48xf32>\n",
      "    %5085 = \"mhlo.dot_general\"(%5084, %5083) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<1x48xf32>, tensor<1xf32>) -> tensor<48xf32>\n",
      "    %5086 = \"mhlo.concatenate\"(%4980, %5001, %5022, %5043, %5064, %5085) {dimension = 0 : i64} : (tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>, tensor<48xf32>) -> tensor<288xf32>\n",
      "    %5087 = \"mhlo.slice\"(%arg10) {limit_indices = dense<[6, 288, 288]> : tensor<3xi64>, start_indices = dense<[5, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x288xf32>) -> tensor<1x288x288xf32>\n",
      "    %5088 = mhlo.reshape %5087 : (tensor<1x288x288xf32>) -> tensor<288x288xf32>\n",
      "    %5089 = \"mhlo.dot_general\"(%5088, %5086) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x288xf32>, tensor<288xf32>) -> tensor<288xf32>\n",
      "    %5090 = mhlo.add %4923, %5089 : tensor<288xf32>\n",
      "    %5091 = \"mhlo.slice\"(%arg2) {limit_indices = dense<[6, 288]> : tensor<2xi64>, start_indices = dense<[5, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<6x288xf32>) -> tensor<1x288xf32>\n",
      "    %5092 = mhlo.reshape %5091 : (tensor<1x288xf32>) -> tensor<288xf32>\n",
      "    %5093 = \"mhlo.dot_general\"(%5090, %5090) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288xf32>, tensor<288xf32>) -> tensor<f32>\n",
      "    %5094 = mhlo.constant dense<2.880000e+02> : tensor<f32>\n",
      "    %5095 = mhlo.divide %5093, %5094 : tensor<f32>\n",
      "    %5096 = mhlo.constant dense<9.99999974E-6> : tensor<f32>\n",
      "    %5097 = mhlo.add %5095, %5096 : tensor<f32>\n",
      "    %5098 = mhlo.sqrt %5097 : tensor<f32>\n",
      "    %5099 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %5100 = mhlo.divide %5099, %5098 : tensor<f32>\n",
      "    %5101 = mhlo.multiply %5092, %5090 : tensor<288xf32>\n",
      "    %5102 = \"mhlo.broadcast_in_dim\"(%5100) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<288xf32>\n",
      "    %5103 = mhlo.multiply %5101, %5102 : tensor<288xf32>\n",
      "    %5104 = \"mhlo.slice\"(%arg5) {limit_indices = dense<[6, 768, 288]> : tensor<3xi64>, start_indices = dense<[5, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x768x288xf32>) -> tensor<1x768x288xf32>\n",
      "    %5105 = mhlo.reshape %5104 : (tensor<1x768x288xf32>) -> tensor<768x288xf32>\n",
      "    %5106 = \"mhlo.dot_general\"(%5105, %5103) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<768x288xf32>, tensor<288xf32>) -> tensor<768xf32>\n",
      "    %5107 = \"mhlo.slice\"(%arg7) {limit_indices = dense<[6, 768, 288]> : tensor<3xi64>, start_indices = dense<[5, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x768x288xf32>) -> tensor<1x768x288xf32>\n",
      "    %5108 = mhlo.reshape %5107 : (tensor<1x768x288xf32>) -> tensor<768x288xf32>\n",
      "    %5109 = \"mhlo.dot_general\"(%5108, %5103) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<768x288xf32>, tensor<288xf32>) -> tensor<768xf32>\n",
      "    %5110 = mhlo.negate %5106 : tensor<768xf32>\n",
      "    %5111 = mhlo.exponential %5110 : tensor<768xf32>\n",
      "    %5112 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %5113 = \"mhlo.broadcast_in_dim\"(%5112) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<768xf32>\n",
      "    %5114 = mhlo.add %5113, %5111 : tensor<768xf32>\n",
      "    %5115 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %5116 = \"mhlo.broadcast_in_dim\"(%5115) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<768xf32>\n",
      "    %5117 = mhlo.divide %5116, %5114 : tensor<768xf32>\n",
      "    %5118 = mhlo.multiply %5106, %5117 : tensor<768xf32>\n",
      "    %5119 = mhlo.multiply %5118, %5109 : tensor<768xf32>\n",
      "    %5120 = \"mhlo.slice\"(%arg6) {limit_indices = dense<[6, 288, 768]> : tensor<3xi64>, start_indices = dense<[5, 0, 0]> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} : (tensor<6x288x768xf32>) -> tensor<1x288x768xf32>\n",
      "    %5121 = mhlo.reshape %5120 : (tensor<1x288x768xf32>) -> tensor<288x768xf32>\n",
      "    %5122 = \"mhlo.dot_general\"(%5121, %5119) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288x768xf32>, tensor<768xf32>) -> tensor<288xf32>\n",
      "    %5123 = mhlo.add %5090, %5122 : tensor<288xf32>\n",
      "    %5124 = \"mhlo.dot_general\"(%5123, %5123) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<288xf32>, tensor<288xf32>) -> tensor<f32>\n",
      "    %5125 = mhlo.constant dense<2.880000e+02> : tensor<f32>\n",
      "    %5126 = mhlo.divide %5124, %5125 : tensor<f32>\n",
      "    %5127 = mhlo.constant dense<9.99999974E-6> : tensor<f32>\n",
      "    %5128 = mhlo.add %5126, %5127 : tensor<f32>\n",
      "    %5129 = mhlo.sqrt %5128 : tensor<f32>\n",
      "    %5130 = mhlo.constant dense<1.000000e+00> : tensor<f32>\n",
      "    %5131 = mhlo.divide %5130, %5129 : tensor<f32>\n",
      "    %5132 = mhlo.multiply %arg3, %5123 : tensor<288xf32>\n",
      "    %5133 = \"mhlo.broadcast_in_dim\"(%5131) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<288xf32>\n",
      "    %5134 = mhlo.multiply %5132, %5133 : tensor<288xf32>\n",
      "    %5135 = \"mhlo.dot_general\"(%arg8, %5134) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]} : (tensor<32000x288xf32>, tensor<288xf32>) -> tensor<32000xf32>\n",
      "    %5136 = \"mhlo.broadcast_in_dim\"(%3955) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x1x288xf32>\n",
      "    %5137 = \"mhlo.broadcast_in_dim\"(%4155) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x1x288xf32>\n",
      "    %5138 = \"mhlo.broadcast_in_dim\"(%4355) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x1x288xf32>\n",
      "    %5139 = \"mhlo.broadcast_in_dim\"(%4555) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x1x288xf32>\n",
      "    %5140 = \"mhlo.broadcast_in_dim\"(%4755) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x1x288xf32>\n",
      "    %5141 = \"mhlo.broadcast_in_dim\"(%4955) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x1x288xf32>\n",
      "    %5142 = \"mhlo.concatenate\"(%5136, %5137, %5138, %5139, %5140, %5141) {dimension = 0 : i64} : (tensor<1x1x288xf32>, tensor<1x1x288xf32>, tensor<1x1x288xf32>, tensor<1x1x288xf32>, tensor<1x1x288xf32>, tensor<1x1x288xf32>) -> tensor<6x1x288xf32>\n",
      "    %5143 = \"mhlo.broadcast_in_dim\"(%3959) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x1x288xf32>\n",
      "    %5144 = \"mhlo.broadcast_in_dim\"(%4159) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x1x288xf32>\n",
      "    %5145 = \"mhlo.broadcast_in_dim\"(%4359) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x1x288xf32>\n",
      "    %5146 = \"mhlo.broadcast_in_dim\"(%4559) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x1x288xf32>\n",
      "    %5147 = \"mhlo.broadcast_in_dim\"(%4759) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x1x288xf32>\n",
      "    %5148 = \"mhlo.broadcast_in_dim\"(%4959) {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} : (tensor<1x288xf32>) -> tensor<1x1x288xf32>\n",
      "    %5149 = \"mhlo.concatenate\"(%5143, %5144, %5145, %5146, %5147, %5148) {dimension = 0 : i64} : (tensor<1x1x288xf32>, tensor<1x1x288xf32>, tensor<1x1x288xf32>, tensor<1x1x288xf32>, tensor<1x1x288xf32>, tensor<1x1x288xf32>) -> tensor<6x1x288xf32>\n",
      "    return %5135, %5142, %5149 : tensor<32000xf32>, tensor<6x1x288xf32>, tensor<6x1x288xf32>\n",
      "  }\n",
      "  func.func private @append(%arg0: tensor<0x288xf32>, %arg1: tensor<1x288xf32>) -> tensor<1x288xf32> {\n",
      "    %0 = \"mhlo.concatenate\"(%arg0, %arg1) {dimension = 0 : i64} : (tensor<0x288xf32>, tensor<1x288xf32>) -> tensor<1x288xf32>\n",
      "    return %0 : tensor<1x288xf32>\n",
      "  }\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(str(mlir))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 220,
   "id": "e8184b32-37eb-41e9-8d69-af983abda471",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.99 ms ± 161 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit jfunc(1, weights, key_cache, value_cache)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6de3fb3c-ad2e-4db6-96eb-be0df98b1d88",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}