API tools faq
paste
Guest User

jaxpr

a guest
Aug 1st, 2022
14
0
Never
Add comment
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 8.75 KB | None | 0 0
raw download clone embed print report
  1. { lambda ; a:f32[3,3] b:f32[3,3]. let
  2.  
  3. c:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
  4. d:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
  5. e:i32[2] = concatenate[dimension=0] c d
  6. f:f32[2,2] = gather[
  7. dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(0, 1))
  8. fill_value=None
  9. indices_are_sorted=True
  10. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  11. slice_sizes=(2, 2)
  12. unique_indices=True
  13. ] a e
  14.  
  15. g:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
  16. h:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 1
  17. i:i32[2] = concatenate[dimension=0] g h
  18. j:f32[2,2] = gather[
  19. dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(0, 1))
  20. fill_value=None
  21. indices_are_sorted=True
  22. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  23. slice_sizes=(2, 2)
  24. unique_indices=True
  25. ] a i
  26.  
  27. k:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 1
  28. l:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
  29. m:i32[2] = concatenate[dimension=0] k l
  30. n:f32[2,2] = gather[
  31. dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(0, 1))
  32. fill_value=None
  33. indices_are_sorted=True
  34. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  35. slice_sizes=(2, 2)
  36. unique_indices=True
  37. ] a m
  38.  
  39. o:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 1
  40. p:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 1
  41. q:i32[2] = concatenate[dimension=0] o p
  42. r:f32[2,2] = gather[
  43. dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(0, 1))
  44. fill_value=None
  45. indices_are_sorted=True
  46. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  47. slice_sizes=(2, 2)
  48. unique_indices=True
  49. ] a q
  50.  
  51. s:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
  52. t:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
  53. u:i32[2] = concatenate[dimension=0] s t
  54. v:f32[2,2] = gather[
  55. dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(0, 1))
  56. fill_value=None
  57. indices_are_sorted=True
  58. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  59. slice_sizes=(2, 2)
  60. unique_indices=True
  61. ] b u
  62.  
  63. w:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
  64. x:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 1
  65. y:i32[2] = concatenate[dimension=0] w x
  66. z:f32[2,2] = gather[
  67. dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(0, 1))
  68. fill_value=None
  69. indices_are_sorted=True
  70. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  71. slice_sizes=(2, 2)
  72. unique_indices=True
  73. ] b y
  74.  
  75. ba:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 1
  76. bb:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
  77. bc:i32[2] = concatenate[dimension=0] ba bb
  78. bd:f32[2,2] = gather[
  79. dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(0, 1))
  80. fill_value=None
  81. indices_are_sorted=True
  82. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  83. slice_sizes=(2, 2)
  84. unique_indices=True
  85. ] b bc
  86.  
  87. be:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 1
  88. bf:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 1
  89. bg:i32[2] = concatenate[dimension=0] be bf
  90. bh:f32[2,2] = gather[
  91. dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(), start_index_map=(0, 1))
  92. fill_value=None
  93. indices_are_sorted=True
  94. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  95. slice_sizes=(2, 2)
  96. unique_indices=True
  97. ] b bg
  98.  
  99. bi:f32[2,2] = add f 0.0
  100. bj:f32[2,2] = add bi j
  101. bk:f32[2,2] = add bj n
  102. bl:f32[2,2] = add bk r
  103. bm:f32[2,2] = add bl v
  104. bn:f32[2,2] = add bm z
  105. bo:f32[2,2] = add bn bd
  106. bp:f32[2,2] = add bo bh
  107. _:f32[] = reduce_sum[axes=(0, 1)] bp
  108. bq:f32[2,2] = broadcast_in_dim[broadcast_dimensions=() shape=(2, 2)] 1.0
  109. br:f32[2,2] = reduce_sum[axes=()] bq
  110. bs:f32[3,3] = broadcast_in_dim[broadcast_dimensions=() shape=(3, 3)] 0.0
  111. bt:f32[3,3] = scatter-add[
  112. dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0, 1), inserted_window_dims=(), scatter_dims_to_operand_dims=(0, 1))
  113. indices_are_sorted=True
  114. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  115. unique_indices=True
  116. update_consts=()
  117. update_jaxpr={ lambda ; bu:f32[] bv:f32[]. let
  118. bw:f32[] = add bu bv
  119. in (bw,) }
  120. ] bs bg br
  121. bx:f32[2,2] = reduce_sum[axes=()] bq
  122. by:f32[3,3] = broadcast_in_dim[broadcast_dimensions=() shape=(3, 3)] 0.0
  123. bz:f32[3,3] = scatter-add[
  124. dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0, 1), inserted_window_dims=(), scatter_dims_to_operand_dims=(0, 1))
  125. indices_are_sorted=True
  126. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  127. unique_indices=True
  128. update_consts=()
  129. update_jaxpr={ lambda ; bu:f32[] bv:f32[]. let
  130. bw:f32[] = add bu bv
  131. in (bw,) }
  132. ] by bc bx
  133. ca:f32[3,3] = add_any bt bz
  134. cb:f32[2,2] = reduce_sum[axes=()] bq
  135. cc:f32[3,3] = broadcast_in_dim[broadcast_dimensions=() shape=(3, 3)] 0.0
  136. cd:f32[3,3] = scatter-add[
  137. dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0, 1), inserted_window_dims=(), scatter_dims_to_operand_dims=(0, 1))
  138. indices_are_sorted=True
  139. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  140. unique_indices=True
  141. update_consts=()
  142. update_jaxpr={ lambda ; bu:f32[] bv:f32[]. let
  143. bw:f32[] = add bu bv
  144. in (bw,) }
  145. ] cc y cb
  146. ce:f32[3,3] = add_any ca cd
  147. cf:f32[2,2] = reduce_sum[axes=()] bq
  148. cg:f32[3,3] = broadcast_in_dim[broadcast_dimensions=() shape=(3, 3)] 0.0
  149. ch:f32[3,3] = scatter-add[
  150. dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0, 1), inserted_window_dims=(), scatter_dims_to_operand_dims=(0, 1))
  151. indices_are_sorted=True
  152. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  153. unique_indices=True
  154. update_consts=()
  155. update_jaxpr={ lambda ; bu:f32[] bv:f32[]. let
  156. bw:f32[] = add bu bv
  157. in (bw,) }
  158. ] cg u cf
  159. ci:f32[3,3] = add_any ce ch
  160. cj:f32[2,2] = reduce_sum[axes=()] bq
  161. ck:f32[3,3] = broadcast_in_dim[broadcast_dimensions=() shape=(3, 3)] 0.0
  162. cl:f32[3,3] = scatter-add[
  163. dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0, 1), inserted_window_dims=(), scatter_dims_to_operand_dims=(0, 1))
  164. indices_are_sorted=True
  165. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  166. unique_indices=True
  167. update_consts=()
  168. update_jaxpr={ lambda ; bu:f32[] bv:f32[]. let
  169. bw:f32[] = add bu bv
  170. in (bw,) }
  171. ] ck q cj
  172. cm:f32[2,2] = reduce_sum[axes=()] bq
  173. cn:f32[3,3] = broadcast_in_dim[broadcast_dimensions=() shape=(3, 3)] 0.0
  174. co:f32[3,3] = scatter-add[
  175. dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0, 1), inserted_window_dims=(), scatter_dims_to_operand_dims=(0, 1))
  176. indices_are_sorted=True
  177. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  178. unique_indices=True
  179. update_consts=()
  180. update_jaxpr={ lambda ; bu:f32[] bv:f32[]. let
  181. bw:f32[] = add bu bv
  182. in (bw,) }
  183. ] cn m cm
  184. cp:f32[3,3] = add_any cl co
  185. cq:f32[2,2] = reduce_sum[axes=()] bq
  186. cr:f32[3,3] = broadcast_in_dim[broadcast_dimensions=() shape=(3, 3)] 0.0
  187. cs:f32[3,3] = scatter-add[
  188. dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0, 1), inserted_window_dims=(), scatter_dims_to_operand_dims=(0, 1))
  189. indices_are_sorted=True
  190. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  191. unique_indices=True
  192. update_consts=()
  193. update_jaxpr={ lambda ; bu:f32[] bv:f32[]. let
  194. bw:f32[] = add bu bv
  195. in (bw,) }
  196. ] cr i cq
  197. ct:f32[3,3] = add_any cp cs
  198. cu:f32[2,2] = reduce_sum[axes=()] bq
  199. cv:f32[3,3] = broadcast_in_dim[broadcast_dimensions=() shape=(3, 3)] 0.0
  200. cw:f32[3,3] = scatter-add[
  201. dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0, 1), inserted_window_dims=(), scatter_dims_to_operand_dims=(0, 1))
  202. indices_are_sorted=True
  203. mode=GatherScatterMode.PROMISE_IN_BOUNDS
  204. unique_indices=True
  205. update_consts=()
  206. update_jaxpr={ lambda ; bu:f32[] bv:f32[]. let
  207. bw:f32[] = add bu bv
  208. in (bw,) }
  209. ] cv e cu
  210. cx:f32[3,3] = add_any ct cw
  211. in (cx, ci) }
Advertisement
Add Comment
Please, Sign In to add comment
Public Pastes
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy.  OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!