Dedup prev_outs and cur_inps in looped_opt_barrier. As you can see below, we get rid of the redundant opt_barriers.

[copybara-service]

Dedup prev_outs and cur_inps in looped_opt_barrier. As you can see below, we get rid of the redundant opt_barriers.

Before jaxpr:

{ lambda ; a:f64[8@x] b:f64[8@x] c:f64[8@x]. let
    d:f64[8@x] = sin a
    e:Tok = create_token
    f:Tok _:f64[8@x] = optimization_barrier e d
    _:Tok g:f64[8@x] = optimization_barrier f d
    h:f64[8@x] = program_order[
      call_jaxpr={ lambda ; i:f64[8@x]. let
          j:f64[8@x] = add i 1.0:f64[]
        in (j,) }
      exclude_mask=(False,)
    ] g
    k:Tok = create_token
    l:Tok _:f64[8@x] = optimization_barrier k h
    _:Tok m:f64[8@x] = optimization_barrier l h
    _:Tok n:f64[8@x] = optimization_barrier l b
    _:Tok o:f64[8@x] = optimization_barrier l c
    p:f64[8@x] = program_order[
      call_jaxpr={ lambda ; q:f64[8@x] r:f64[8@x] s:f64[8@x]. let
          t:f64[8@x] = mul q r
          u:f64[8@x] = add t s
        in (u,) }
      exclude_mask=(False, False, False)
    ] m n o
  in (p,) }

After jaxpr:

{ lambda ; a:f64[8@x] b:f64[8@x] c:f64[8@x]. let
    d:f64[8@x] = sin a
    e:Tok = create_token
    _:Tok f:f64[8@x] = optimization_barrier e d
    g:f64[8@x] = program_order[
      call_jaxpr={ lambda ; h:f64[8@x]. let
          i:f64[8@x] = add h 1.0:f64[]
        in (i,) }
      exclude_mask=(False,)
    ] f
    j:Tok = create_token
    k:Tok l:f64[8@x] = optimization_barrier j g
    _:Tok m:f64[8@x] = optimization_barrier k b
    _:Tok n:f64[8@x] = optimization_barrier k c
    o:f64[8@x] = program_order[
      call_jaxpr={ lambda ; p:f64[8@x] q:f64[8@x] r:f64[8@x]. let
          s:f64[8@x] = mul p q
          t:f64[8@x] = add s r
        in (t,) }
      exclude_mask=(False, False, False)
    ] l m n
  in (o,) }