Add files using upload-large-folder tool

.gitattributes

@@ -4510,3 +4510,5 @@ external/alphageometry/.venv-ag/Lib/site-packages/pip/_vendor/distlib/t64-arm.ex
 external/alphageometry/.venv-ag/Lib/site-packages/pip/_vendor/distlib/t64.exe filter=lfs diff=lfs merge=lfs -text
 external/alphageometry/.venv-ag/Lib/site-packages/pip/_vendor/distlib/w64-arm.exe filter=lfs diff=lfs merge=lfs -text
 external/alphageometry/.venv-ag/Lib/site-packages/pip/_vendor/distlib/w64.exe filter=lfs diff=lfs merge=lfs -text
+hfenv/Lib/site-packages/setuptools/cli-arm64.exe filter=lfs diff=lfs merge=lfs -text
+hfenv/Lib/site-packages/setuptools/gui-arm64.exe filter=lfs diff=lfs merge=lfs -text

backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Elab/Deriving/BEq.lean

@@ -0,0 +1,155 @@
+/-
+Copyright (c) 2020 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import Lean.Meta.Transform
+import Lean.Elab.Deriving.Basic
+import Lean.Elab.Deriving.Util
+
+namespace Lean.Elab.Deriving.BEq
+open Lean.Parser.Term
+open Meta
+
+def mkBEqHeader (indVal : InductiveVal) : TermElabM Header := do
+  mkHeader `BEq 2 indVal
+
+def mkMatch (header : Header) (indVal : InductiveVal) (auxFunName : Name) : TermElabM Term := do
+  let discrs ← mkDiscrs header indVal
+  let alts ← mkAlts
+  `(match $[$discrs],* with $alts:matchAlt*)
+where
+  mkElseAlt : TermElabM (TSyntax ``matchAltExpr) := do
+    let mut patterns := #[]
+    -- add `_` pattern for indices
+    for _ in [:indVal.numIndices] do
+      patterns := patterns.push (← `(_))
+    patterns := patterns.push (← `(_))
+    patterns := patterns.push (← `(_))
+    let altRhs ← `(false)
+    `(matchAltExpr| | $[$patterns:term],* => $altRhs:term)
+
+  mkAlts : TermElabM (Array (TSyntax ``matchAlt)) := do
+    let mut alts := #[]
+    for ctorName in indVal.ctors do
+      let ctorInfo ← getConstInfoCtor ctorName
+      let alt ← forallTelescopeReducing ctorInfo.type fun xs type => do
+        let type ← Core.betaReduce type -- we 'beta-reduce' to eliminate "artificial" dependencies
+        let mut patterns := #[]
+        -- add `_` pattern for indices
+        for _ in [:indVal.numIndices] do
+          patterns := patterns.push (← `(_))
+        let mut ctorArgs1 := #[]
+        let mut ctorArgs2 := #[]
+        let mut rhs ← `(true)
+        let mut rhs_empty := true
+        for i in [:ctorInfo.numFields] do
+          let pos := indVal.numParams + ctorInfo.numFields - i - 1
+          let x := xs[pos]!
+          if type.containsFVar x.fvarId! then
+            -- If resulting type depends on this field, we don't need to compare
+            ctorArgs1 := ctorArgs1.push (← `(_))
+            ctorArgs2 := ctorArgs2.push (← `(_))
+          else
+            let a := mkIdent (← mkFreshUserName `a)
+            let b := mkIdent (← mkFreshUserName `b)
+            ctorArgs1 := ctorArgs1.push a
+            ctorArgs2 := ctorArgs2.push b
+            let xType ← inferType x
+            if (← isProp xType) then
+              continue
+            if xType.isAppOf indVal.name then
+              if rhs_empty then
+                rhs ← `($(mkIdent auxFunName):ident $a:ident $b:ident)
+                rhs_empty := false
+              else
+                rhs ← `($(mkIdent auxFunName):ident $a:ident $b:ident && $rhs)
+            /- If `x` appears in the type of another field, use `eq_of_beq` to
+               unify the types of the subsequent variables -/
+            else if ← xs[(pos+1)...*].anyM
+                (fun fvar => (Expr.containsFVar · x.fvarId!) <$> (inferType fvar)) then
+              rhs ← `(if h : $a:ident == $b:ident then by
+                        cases (eq_of_beq h)
+                        exact $rhs
+                      else false)
+              rhs_empty := false
+            else
+              if rhs_empty then
+                rhs ← `($a:ident == $b:ident)
+                rhs_empty := false
+              else
+                rhs ← `($a:ident == $b:ident && $rhs)
+          -- add `_` for inductive parameters, they are inaccessible
+        for _ in [:indVal.numParams] do
+          ctorArgs1 := ctorArgs1.push (← `(_))
+          ctorArgs2 := ctorArgs2.push (← `(_))
+        patterns := patterns.push (← `(@$(mkIdent ctorName):ident $ctorArgs1.reverse:term*))
+        patterns := patterns.push (← `(@$(mkIdent ctorName):ident $ctorArgs2.reverse:term*))
+        `(matchAltExpr| | $[$patterns:term],* => $rhs:term)
+      alts := alts.push alt
+    alts := alts.push (← mkElseAlt)
+    return alts
+
+def mkAuxFunction (ctx : Context) (i : Nat) : TermElabM Command := do
+  let auxFunName := ctx.auxFunNames[i]!
+  let indVal     := ctx.typeInfos[i]!
+  let header     ← mkBEqHeader indVal
+  let mut body   ← mkMatch header indVal auxFunName
+  if ctx.usePartial then
+    let letDecls ← mkLocalInstanceLetDecls ctx `BEq header.argNames
+    body ← mkLet letDecls body
+  let binders    := header.binders
+  if ctx.usePartial then
+    `(partial def $(mkIdent auxFunName):ident $binders:bracketedBinder* : Bool := $body:term)
+  else
+    `(@[expose] def $(mkIdent auxFunName):ident $binders:bracketedBinder* : Bool := $body:term)
+
+def mkMutualBlock (ctx : Context) : TermElabM Syntax := do
+  let mut auxDefs := #[]
+  for i in [:ctx.typeInfos.size] do
+    auxDefs := auxDefs.push (← mkAuxFunction ctx i)
+  `(mutual
+     set_option match.ignoreUnusedAlts true
+     $auxDefs:command*
+    end)
+
+private def mkBEqInstanceCmds (declName : Name) : TermElabM (Array Syntax) := do
+  let ctx ← mkContext "beq" declName
+  let cmds := #[← mkMutualBlock ctx] ++ (← mkInstanceCmds ctx `BEq #[declName])
+  trace[Elab.Deriving.beq] "\n{cmds}"
+  return cmds
+
+private def mkBEqEnumFun (ctx : Context) (name : Name) : TermElabM Syntax := do
+  let auxFunName := ctx.auxFunNames[0]!
+  `(@[expose] def $(mkIdent auxFunName):ident  (x y : $(mkIdent name)) : Bool := x.toCtorIdx == y.toCtorIdx)
+
+private def mkBEqEnumCmd (name : Name): TermElabM (Array Syntax) := do
+  let ctx ← mkContext "beq" name
+  let cmds := #[← mkBEqEnumFun ctx name] ++ (← mkInstanceCmds ctx `BEq #[name])
+  trace[Elab.Deriving.beq] "\n{cmds}"
+  return cmds
+
+open Command
+
+def mkBEqInstance (declName : Name) : CommandElabM Unit := do
+    let cmds ← liftTermElabM <|
+      if (← isEnumType declName) then
+        mkBEqEnumCmd declName
+      else
+         mkBEqInstanceCmds declName
+    cmds.forM elabCommand
+
+def mkBEqInstanceHandler (declNames : Array Name) : CommandElabM Bool := do
+  if (← declNames.allM isInductive) then
+    for declName in declNames do
+      mkBEqInstance declName
+    return true
+  else
+    return false
+
+builtin_initialize
+  registerDerivingHandler `BEq mkBEqInstanceHandler
+  registerTraceClass `Elab.Deriving.beq
+
+end Lean.Elab.Deriving.BEq

backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Elab/Deriving/Basic.lean

@@ -0,0 +1,134 @@
+/-
+Copyright (c) 2020 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura, Wojciech Nawrocki
+-/
+prelude
+import Lean.Elab.Command
+import Lean.Elab.DeclarationRange
+
+namespace Lean.Elab
+open Command
+
+namespace Term
+open Meta
+
+/-- Result for `mkInst?` -/
+structure MkInstResult where
+  instVal   : Expr
+  instType  : Expr
+  outParams : Array Expr := #[]
+
+/--
+  Construct an instance for `className out₁ ... outₙ type`.
+  The method support classes with a prefix of `outParam`s (e.g. `MonadReader`). -/
+private partial def mkInst? (className : Name) (type : Expr) : MetaM (Option MkInstResult) := do
+  let rec go? (instType instTypeType : Expr) (outParams : Array Expr) : MetaM (Option MkInstResult) := do
+    let instTypeType ← whnfD instTypeType
+    unless instTypeType.isForall do
+      return none
+    let d := instTypeType.bindingDomain!
+    if d.isOutParam then
+      let mvar ← mkFreshExprMVar d
+      go? (mkApp instType mvar) (instTypeType.bindingBody!.instantiate1 mvar) (outParams.push mvar)
+    else
+      unless (← isDefEqGuarded (← inferType type) d) do
+        return none
+      let instType ← instantiateMVars (mkApp instType type)
+      let instVal ← synthInstance instType
+      return some { instVal, instType, outParams }
+  let instType ← mkConstWithFreshMVarLevels className
+  go? instType (← inferType instType) #[]
+
+def processDefDeriving (className : Name) (declName : Name) : TermElabM Bool := do
+  try
+    let ConstantInfo.defnInfo info ← getConstInfo declName | return false
+    let some result ← mkInst? className info.value | return false
+    let instTypeNew := mkApp result.instType.appFn! (Lean.mkConst declName (info.levelParams.map mkLevelParam))
+    Meta.check instTypeNew
+    let instName ← liftMacroM <| mkUnusedBaseName (declName.appendBefore "inst" |>.appendAfter className.getString!)
+    addAndCompile <| Declaration.defnDecl {
+      name        := instName
+      levelParams := info.levelParams
+      type        := (← instantiateMVars instTypeNew)
+      value       := (← instantiateMVars result.instVal)
+      hints       := info.hints
+      safety      := info.safety
+    }
+    addInstance instName AttributeKind.global (eval_prio default)
+    addDeclarationRangesFromSyntax instName (← getRef)
+    return true
+  catch _ =>
+    return false
+
+end Term
+
+def DerivingHandler := (typeNames : Array Name) → CommandElabM Bool
+
+builtin_initialize derivingHandlersRef : IO.Ref (NameMap (List DerivingHandler)) ← IO.mkRef {}
+
+/-- A `DerivingHandler` is called on the fully qualified names of all types it is running for
+as well as the syntax of a `with` argument, if present.
+
+For example, `deriving instance Foo with fooArgs for Bar, Baz` invokes
+``fooHandler #[`Bar, `Baz] `(fooArgs)``. -/
+def registerDerivingHandler (className : Name) (handler : DerivingHandler) : IO Unit := do
+  unless (← initializing) do
+    throw (IO.userError "failed to register deriving handler, it can only be registered during initialization")
+  derivingHandlersRef.modify fun m => match m.find? className with
+    | some handlers => m.insert className (handler :: handlers)
+    | none => m.insert className [handler]
+
+def defaultHandler (className : Name) (typeNames : Array Name) : CommandElabM Unit := do
+  throwError "default handlers have not been implemented yet, class: '{className}' types: {typeNames}"
+
+def applyDerivingHandlers (className : Name) (typeNames : Array Name) : CommandElabM Unit := do
+  withTraceNode `Elab.Deriving (fun _ => return m!"running deriving handlers for '{className}'") do
+    match (← derivingHandlersRef.get).find? className with
+    | some handlers =>
+      for handler in handlers do
+        if (← handler typeNames) then
+          return ()
+      defaultHandler className typeNames
+    | none => defaultHandler className typeNames
+
+private def tryApplyDefHandler (className : Name) (declName : Name) : CommandElabM Bool :=
+  liftTermElabM do
+    Term.processDefDeriving className declName
+
+@[builtin_command_elab «deriving»] def elabDeriving : CommandElab
+  | `(deriving instance $[$classes],* for $[$declNames],*) => do
+     let declNames ← liftCoreM <| declNames.mapM realizeGlobalConstNoOverloadWithInfo
+     for cls in classes do
+       try
+         let className ← liftCoreM <| realizeGlobalConstNoOverloadWithInfo cls
+         withRef cls do
+           if declNames.size == 1 then
+             if (← tryApplyDefHandler className declNames[0]!) then
+               return ()
+           applyDerivingHandlers className declNames
+       catch ex =>
+         logException ex
+  | _ => throwUnsupportedSyntax
+
+structure DerivingClassView where
+  ref : Syntax
+  className : Name
+
+def getOptDerivingClasses (optDeriving : Syntax) : CoreM (Array DerivingClassView) := do
+  match optDeriving with
+  | `(Parser.Command.optDeriving| deriving $[$classes],*) =>
+    let mut ret := #[]
+    for cls in classes do
+      let className ← realizeGlobalConstNoOverloadWithInfo cls
+      ret := ret.push { ref := cls, className := className }
+    return ret
+  | _ => return #[]
+
+def DerivingClassView.applyHandlers (view : DerivingClassView) (declNames : Array Name) : CommandElabM Unit :=
+  withRef view.ref do applyDerivingHandlers view.className declNames
+
+builtin_initialize
+  registerTraceClass `Elab.Deriving
+
+end Lean.Elab

backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Elab/Deriving/DecEq.lean

@@ -0,0 +1,212 @@
+/-
+Copyright (c) 2020 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import Lean.Meta.Transform
+import Lean.Meta.Inductive
+import Lean.Elab.Deriving.Basic
+import Lean.Elab.Deriving.Util
+
+namespace Lean.Elab.Deriving.DecEq
+open Lean.Parser.Term
+open Meta
+
+def mkDecEqHeader (indVal : InductiveVal) : TermElabM Header := do
+  mkHeader `DecidableEq 2 indVal
+
+def mkMatch (ctx : Context) (header : Header) (indVal : InductiveVal) : TermElabM Term := do
+  let discrs ← mkDiscrs header indVal
+  let alts ← mkAlts
+  `(match $[$discrs],* with $alts:matchAlt*)
+where
+  mkSameCtorRhs : List (Ident × Ident × Option Name × Bool) → TermElabM Term
+    | [] => ``(isTrue rfl)
+    | (a, b, recField, isProof) :: todo => withFreshMacroScope do
+      let rhs ← if isProof then
+        `(have h : @$a = @$b := rfl; by subst h; exact $(← mkSameCtorRhs todo):term)
+      else
+        let sameCtor ← mkSameCtorRhs todo
+        `(if h : @$a = @$b then
+           by subst h; exact $sameCtor:term
+          else
+           isFalse (by intro n; injection n; apply h _; assumption))
+      if let some auxFunName := recField then
+        -- add local instance for `a = b` using the function being defined `auxFunName`
+        `(let inst := $(mkIdent auxFunName) @$a @$b; $rhs)
+      else
+        return rhs
+
+  mkAlts : TermElabM (Array (TSyntax ``matchAlt)) := do
+    let mut alts := #[]
+    for ctorName₁ in indVal.ctors do
+      let ctorInfo ← getConstInfoCtor ctorName₁
+      for ctorName₂ in indVal.ctors do
+        let mut patterns := #[]
+        -- add `_` pattern for indices
+        for _ in [:indVal.numIndices] do
+          patterns := patterns.push (← `(_))
+        if ctorName₁ == ctorName₂ then
+          let alt ← forallTelescopeReducing ctorInfo.type fun xs type => do
+            let type ← Core.betaReduce type -- we 'beta-reduce' to eliminate "artificial" dependencies
+            let mut patterns  := patterns
+            let mut ctorArgs1 := #[]
+            let mut ctorArgs2 := #[]
+            -- add `_` for inductive parameters, they are inaccessible
+            for _ in [:indVal.numParams] do
+              ctorArgs1 := ctorArgs1.push (← `(_))
+              ctorArgs2 := ctorArgs2.push (← `(_))
+            let mut todo := #[]
+            for i in [:ctorInfo.numFields] do
+              let x := xs[indVal.numParams + i]!
+              if type.containsFVar x.fvarId! then
+                -- If resulting type depends on this field, we don't need to compare
+                ctorArgs1 := ctorArgs1.push (← `(_))
+                ctorArgs2 := ctorArgs2.push (← `(_))
+              else
+                let a := mkIdent (← mkFreshUserName `a)
+                let b := mkIdent (← mkFreshUserName `b)
+                ctorArgs1 := ctorArgs1.push a
+                ctorArgs2 := ctorArgs2.push b
+                let xType ← inferType x
+                let indValNum :=
+                  ctx.typeInfos.findIdx?
+                  (xType.isAppOf ∘ ConstantVal.name ∘ InductiveVal.toConstantVal)
+                let recField  := indValNum.map (ctx.auxFunNames[·]!)
+                let isProof ← isProp xType
+                todo := todo.push (a, b, recField, isProof)
+            patterns := patterns.push (← `(@$(mkIdent ctorName₁):ident $ctorArgs1:term*))
+            patterns := patterns.push (← `(@$(mkIdent ctorName₁):ident $ctorArgs2:term*))
+            let rhs ← mkSameCtorRhs todo.toList
+            `(matchAltExpr| | $[$patterns:term],* => $rhs:term)
+          alts := alts.push alt
+        else if (← compatibleCtors ctorName₁ ctorName₂) then
+          patterns := patterns ++ #[(← `($(mkIdent ctorName₁) ..)), (← `($(mkIdent ctorName₂) ..))]
+          let rhs ← `(isFalse (by intro h; injection h))
+          alts := alts.push (← `(matchAltExpr| | $[$patterns:term],* => $rhs:term))
+    return alts
+
+def mkAuxFunction (ctx : Context) (auxFunName : Name) (indVal : InductiveVal): TermElabM (TSyntax `command) := do
+  let header  ← mkDecEqHeader indVal
+  let body    ← mkMatch ctx header indVal
+  let binders := header.binders
+  let target₁ := mkIdent header.targetNames[0]!
+  let target₂ := mkIdent header.targetNames[1]!
+  let termSuffix ← if indVal.isRec
+    then `(Parser.Termination.suffix|termination_by structural $target₁)
+    else `(Parser.Termination.suffix|)
+  let type    ← `(Decidable ($target₁ = $target₂))
+  `(def $(mkIdent auxFunName):ident $binders:bracketedBinder* : $type:term := $body:term
+    $termSuffix:suffix)
+
+def mkAuxFunctions (ctx : Context) : TermElabM (TSyntax `command) := do
+  let mut res : Array (TSyntax `command) := #[]
+  for i in [:ctx.auxFunNames.size] do
+    let auxFunName := ctx.auxFunNames[i]!
+    let indVal     := ctx.typeInfos[i]!
+    res := res.push (← mkAuxFunction ctx auxFunName indVal)
+  `(command| mutual $[$res:command]* end)
+
+def mkDecEqCmds (indVal : InductiveVal) : TermElabM (Array Syntax) := do
+  let ctx ← mkContext "decEq" indVal.name
+  let cmds := #[← mkAuxFunctions ctx] ++ (← mkInstanceCmds ctx `DecidableEq #[indVal.name] (useAnonCtor := false))
+  trace[Elab.Deriving.decEq] "\n{cmds}"
+  return cmds
+
+open Command
+
+def mkDecEq (declName : Name) : CommandElabM Bool := do
+  let indVal ← getConstInfoInduct declName
+  if indVal.isNested then
+    return false -- nested inductive types are not supported yet
+  else
+    let cmds ← liftTermElabM <| mkDecEqCmds indVal
+    -- `cmds` can have a number of syntax nodes quadratic in the number of constructors
+    -- and thus create as many info tree nodes, which we never make use of but which can
+    -- significantly slow down e.g. the unused variables linter; avoid creating them
+    withEnableInfoTree false do
+      cmds.forM elabCommand
+    return true
+
+partial def mkEnumOfNat (declName : Name) : MetaM Unit := do
+  let indVal ← getConstInfoInduct declName
+  let enumType := mkConst declName
+  let ctors := indVal.ctors.toArray
+  withLocalDeclD `n (mkConst ``Nat) fun n => do
+    let cond := mkConst ``cond [1]
+    let rec mkDecTree (low high : Nat) : Expr :=
+      if low + 1 == high then
+        mkConst ctors[low]!
+      else if low + 2 == high then
+        mkApp4 cond enumType (mkApp2 (mkConst ``Nat.beq) n (mkRawNatLit low)) (mkConst ctors[low]!) (mkConst ctors[low+1]!)
+      else
+        let mid := (low + high)/2
+        let lowBranch := mkDecTree low mid
+        let highBranch := mkDecTree mid high
+        mkApp4 cond enumType (mkApp2 (mkConst ``Nat.ble) (mkRawNatLit mid) n) highBranch lowBranch
+    let value ← mkLambdaFVars #[n] (mkDecTree 0 ctors.size)
+    let type ← mkArrow (mkConst ``Nat) enumType
+    addAndCompile <| Declaration.defnDecl {
+      name := Name.mkStr declName "ofNat"
+      levelParams := []
+      safety := DefinitionSafety.safe
+      hints  := ReducibilityHints.abbrev
+      value, type
+    }
+
+def mkEnumOfNatThm (declName : Name) : MetaM Unit := do
+  let indVal ← getConstInfoInduct declName
+  let toCtorIdx := mkConst (Name.mkStr declName "toCtorIdx")
+  let ofNat     := mkConst (Name.mkStr declName "ofNat")
+  let enumType  := mkConst declName
+  let eqEnum    := mkApp (mkConst ``Eq [levelOne]) enumType
+  let rflEnum   := mkApp (mkConst ``Eq.refl [levelOne]) enumType
+  let ctors := indVal.ctors
+  withLocalDeclD `x enumType fun x => do
+    let resultType := mkApp2 eqEnum (mkApp ofNat (mkApp toCtorIdx x)) x
+    let motive     ← mkLambdaFVars #[x] resultType
+    let casesOn    := mkConst (mkCasesOnName declName) [levelZero]
+    let mut value  := mkApp2 casesOn motive x
+    for ctor in ctors do
+      value := mkApp value (mkApp rflEnum (mkConst ctor))
+    value ← mkLambdaFVars #[x] value
+    let type ← mkForallFVars #[x] resultType
+    addAndCompile <| Declaration.thmDecl {
+      name := Name.mkStr declName "ofNat_toCtorIdx"
+      levelParams := []
+      value, type
+    }
+
+def mkDecEqEnum (declName : Name) : CommandElabM Unit := do
+  liftTermElabM <| mkEnumOfNat declName
+  liftTermElabM <| mkEnumOfNatThm declName
+  let ofNatIdent  := mkIdent (Name.mkStr declName "ofNat")
+  let auxThmIdent := mkIdent (Name.mkStr declName "ofNat_toCtorIdx")
+  let cmd ← `(
+    instance : DecidableEq $(mkIdent declName) :=
+      fun x y =>
+        if h : x.toCtorIdx = y.toCtorIdx then
+          -- We use `rfl` in the following proof because the first script fails for unit-like datatypes due to etaStruct.
+          isTrue (by first | have aux := congrArg $ofNatIdent h; rw [$auxThmIdent:ident, $auxThmIdent:ident] at aux; assumption | rfl)
+        else
+          isFalse fun h => by subst h; contradiction
+  )
+  trace[Elab.Deriving.decEq] "\n{cmd}"
+  elabCommand cmd
+
+def mkDecEqInstance (declName : Name) : CommandElabM Bool := do
+  if (← isEnumType declName) then
+    mkDecEqEnum declName
+    return true
+  else
+    mkDecEq declName
+
+def mkDecEqInstanceHandler (declNames : Array Name) : CommandElabM Bool := do
+  declNames.foldlM (fun b n => andM (pure b) (mkDecEqInstance n)) true
+
+builtin_initialize
+  registerDerivingHandler `DecidableEq mkDecEqInstanceHandler
+  registerTraceClass `Elab.Deriving.decEq
+
+end Lean.Elab.Deriving.DecEq

backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Elab/Deriving/FromToJson.lean

@@ -0,0 +1,249 @@
+/-
+Copyright (c) 2020 Sebastian Ullrich. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Sebastian Ullrich, Dany Fabian
+-/
+prelude
+import Lean.Meta.Transform
+import Lean.Elab.Deriving.Basic
+import Lean.Elab.Deriving.Util
+import Lean.Data.Json.FromToJson
+
+namespace Lean.Elab.Deriving.FromToJson
+open Lean.Elab.Command
+open Lean.Json
+open Lean.Parser.Term
+open Lean.Meta
+
+def mkToJsonHeader (indVal : InductiveVal) : TermElabM Header := do
+  mkHeader ``ToJson 1 indVal
+
+def mkFromJsonHeader (indVal : InductiveVal) : TermElabM Header := do
+  let header ← mkHeader ``FromJson 0 indVal
+  let jsonArg ← `(bracketedBinderF|(json : Json))
+  return {header with
+    binders := header.binders.push jsonArg}
+
+def mkJsonField (n : Name) : CoreM (Bool × Term) := do
+  let .str .anonymous s := n | throwError "invalid json field name {n}"
+  let s₁ := s.dropRightWhile (· == '?')
+  return (s != s₁, Syntax.mkStrLit s₁)
+
+def mkToJsonBodyForStruct (header : Header) (indName : Name) : TermElabM Term := do
+  let fields := getStructureFieldsFlattened (← getEnv) indName (includeSubobjectFields := false)
+  let fields ← fields.mapM fun field => do
+    let (isOptField, nm) ← mkJsonField field
+    let target := mkIdent header.targetNames[0]!
+    if isOptField then ``(opt $nm $target.$(mkIdent field))
+    else ``([($nm, toJson ($target).$(mkIdent field))])
+  `(mkObj <| List.flatten [$fields,*])
+
+def mkToJsonBodyForInduct (ctx : Context) (header : Header) (indName : Name) : TermElabM Term := do
+  let indVal ← getConstInfoInduct indName
+  let toJsonFuncId := mkIdent ctx.auxFunNames[0]!
+  -- Return syntax to JSONify `id`, either via `ToJson` or recursively
+  -- if `id`'s type is the type we're deriving for.
+  let mkToJson (id : Ident) (type : Expr) : TermElabM Term := do
+        if type.isAppOf indVal.name then `($toJsonFuncId:ident $id:ident)
+        else ``(toJson $id:ident)
+  let discrs ← mkDiscrs header indVal
+  let alts ← mkAlts indVal fun ctor args userNames => do
+    let ctorStr := ctor.name.eraseMacroScopes.getString!
+    match args, userNames with
+    | #[], _ => ``(toJson $(quote ctorStr))
+    | #[(x, t)], none => ``(mkObj [($(quote ctorStr), $(← mkToJson x t))])
+    | xs, none =>
+      let xs ← xs.mapM fun (x, t) => mkToJson x t
+      ``(mkObj [($(quote ctorStr), Json.arr #[$[$xs:term],*])])
+    | xs, some userNames =>
+      let xs ← xs.mapIdxM fun idx (x, t) => do
+        `(($(quote userNames[idx]!.eraseMacroScopes.getString!), $(← mkToJson x t)))
+      ``(mkObj [($(quote ctorStr), mkObj [$[$xs:term],*])])
+  `(match $[$discrs],* with $alts:matchAlt*)
+
+where
+  mkAlts
+    (indVal : InductiveVal)
+    (rhs : ConstructorVal → Array (Ident × Expr) → Option (Array Name) → TermElabM Term): TermElabM (Array (TSyntax ``matchAlt)) := do
+      let mut alts := #[]
+      for ctorName in indVal.ctors do
+        let ctorInfo ← getConstInfoCtor ctorName
+        let alt ← forallTelescopeReducing ctorInfo.type fun xs _ => do
+          let mut patterns := #[]
+          -- add `_` pattern for indices
+          for _ in [:indVal.numIndices] do
+            patterns := patterns.push (← `(_))
+          let mut ctorArgs := #[]
+          -- add `_` for inductive parameters, they are inaccessible
+          for _ in [:indVal.numParams] do
+            ctorArgs := ctorArgs.push (← `(_))
+          -- bound constructor arguments and their types
+          let mut binders := #[]
+          let mut userNames := #[]
+          for i in [:ctorInfo.numFields] do
+            let x := xs[indVal.numParams + i]!
+            let localDecl ← x.fvarId!.getDecl
+            if !localDecl.userName.hasMacroScopes then
+              userNames := userNames.push localDecl.userName
+            let a := mkIdent (← mkFreshUserName `a)
+            binders := binders.push (a, localDecl.type)
+            ctorArgs := ctorArgs.push a
+          patterns := patterns.push (← `(@$(mkIdent ctorInfo.name):ident $ctorArgs:term*))
+          let rhs ← rhs ctorInfo binders (if userNames.size == binders.size then some userNames else none)
+          `(matchAltExpr| | $[$patterns:term],* => $rhs:term)
+        alts := alts.push alt
+      return alts
+
+def mkFromJsonBodyForStruct (indName : Name) : TermElabM Term := do
+  let fields := getStructureFieldsFlattened (← getEnv) indName (includeSubobjectFields := false)
+  let getters ← fields.mapM (fun field => do
+    let getter ← `(getObjValAs? json _ $(Prod.snd <| ← mkJsonField field))
+    let getter ← `(doElem| Except.mapError (fun s => (toString $(quote indName)) ++ "." ++ (toString $(quote field)) ++ ": " ++ s) <| $getter)
+    return getter
+  )
+  let fields := fields.map mkIdent
+  `(do
+    $[let $fields:ident ← $getters]*
+    return { $[$fields:ident := $(id fields)],* })
+
+def mkFromJsonBodyForInduct (ctx : Context) (indName : Name) : TermElabM Term := do
+  let indVal ← getConstInfoInduct indName
+  let alts ← mkAlts indVal
+  let auxTerm ← alts.foldrM (fun xs x => `(Except.orElseLazy $xs (fun _ => $x))) (← `(Except.error "no inductive constructor matched"))
+  `($auxTerm)
+where
+  mkAlts (indVal : InductiveVal) : TermElabM (Array Term) := do
+  let mut alts := #[]
+  for ctorName in indVal.ctors do
+    let ctorInfo ← getConstInfoCtor ctorName
+    let alt ← do forallTelescopeReducing ctorInfo.type fun xs _ => do
+        let mut binders   := #[]
+        let mut userNames := #[]
+        for i in [:ctorInfo.numFields] do
+          let x := xs[indVal.numParams + i]!
+          let localDecl ← x.fvarId!.getDecl
+          if !localDecl.userName.hasMacroScopes then
+            userNames := userNames.push localDecl.userName
+          let a := mkIdent (← mkFreshUserName `a)
+          binders := binders.push (a, localDecl.type)
+        let fromJsonFuncId := mkIdent ctx.auxFunNames[0]!
+        -- Return syntax to parse `id`, either via `FromJson` or recursively
+        -- if `id`'s type is the type we're deriving for.
+        let mkFromJson (idx : Nat) (type : Expr) : TermElabM (TSyntax ``doExpr) :=
+          if type.isAppOf indVal.name then `(Lean.Parser.Term.doExpr| $fromJsonFuncId:ident jsons[$(quote idx)]!)
+          else `(Lean.Parser.Term.doExpr| fromJson? jsons[$(quote idx)]!)
+        let identNames := binders.map Prod.fst
+        let fromJsons ← binders.mapIdxM fun idx (_, type) => mkFromJson idx type
+        let userNamesOpt ← if binders.size == userNames.size then
+          ``(some #[$[$(userNames.map quote)],*])
+        else
+          ``(none)
+        let stx ←
+          `((Json.parseTagged json $(quote ctorName.eraseMacroScopes.getString!) $(quote ctorInfo.numFields) $(quote userNamesOpt)).bind
+            (fun jsons => do
+              $[let $identNames:ident ← $fromJsons:doExpr]*
+              return $(mkIdent ctorName):ident $identNames*))
+        pure (stx, ctorInfo.numFields)
+      alts := alts.push alt
+  -- the smaller cases, especially the ones without fields are likely faster
+  let alts' := alts.qsort (fun (_, x) (_, y) => x < y)
+  return alts'.map Prod.fst
+
+def mkToJsonBody (ctx : Context) (header : Header) (e : Expr): TermElabM Term := do
+  let indName := e.getAppFn.constName!
+  if isStructure (← getEnv) indName then
+    mkToJsonBodyForStruct header indName
+  else
+    mkToJsonBodyForInduct ctx header indName
+
+def mkToJsonAuxFunction (ctx : Context) (i : Nat) : TermElabM Command := do
+  let auxFunName := ctx.auxFunNames[i]!
+  let header     ←  mkToJsonHeader ctx.typeInfos[i]!
+  let binders    := header.binders
+  Term.elabBinders binders fun _ => do
+  let type       ← Term.elabTerm header.targetType none
+  let mut body   ← mkToJsonBody ctx header type
+  if ctx.usePartial then
+    let letDecls ← mkLocalInstanceLetDecls ctx ``ToJson header.argNames
+    body ← mkLet letDecls body
+    `(private partial def $(mkIdent auxFunName):ident $binders:bracketedBinder* : Json := $body:term)
+  else
+    `(private def $(mkIdent auxFunName):ident $binders:bracketedBinder* : Json := $body:term)
+
+def mkFromJsonBody (ctx : Context) (e : Expr) : TermElabM Term := do
+  let indName := e.getAppFn.constName!
+  if isStructure (← getEnv) indName then
+    mkFromJsonBodyForStruct indName
+  else
+    mkFromJsonBodyForInduct ctx indName
+
+def mkFromJsonAuxFunction (ctx : Context) (i : Nat) : TermElabM Command := do
+  let auxFunName := ctx.auxFunNames[i]!
+  let indval     := ctx.typeInfos[i]!
+  let header     ←  mkFromJsonHeader indval --TODO fix header info
+  let binders    := header.binders
+  Term.elabBinders binders fun _ => do
+  let type ← Term.elabTerm header.targetType none
+  let mut body ← mkFromJsonBody ctx type
+  if ctx.usePartial || indval.isRec then
+    let letDecls ← mkLocalInstanceLetDecls ctx ``FromJson header.argNames
+    body ← mkLet letDecls body
+    `(private partial def $(mkIdent auxFunName):ident $binders:bracketedBinder* : Except String $(← mkInductiveApp ctx.typeInfos[i]! header.argNames) := $body:term)
+  else
+    `(private def $(mkIdent auxFunName):ident $binders:bracketedBinder* : Except String $(← mkInductiveApp ctx.typeInfos[i]! header.argNames) := $body:term)
+
+
+def mkToJsonMutualBlock (ctx : Context) : TermElabM Command := do
+  let mut auxDefs := #[]
+  for i in [:ctx.typeInfos.size] do
+    auxDefs := auxDefs.push (← mkToJsonAuxFunction ctx i)
+  `(mutual
+     $auxDefs:command*
+    end)
+
+def mkFromJsonMutualBlock (ctx : Context) : TermElabM Command := do
+  let mut auxDefs := #[]
+  for i in [:ctx.typeInfos.size] do
+    auxDefs := auxDefs.push (← mkFromJsonAuxFunction ctx i)
+  `(mutual
+     $auxDefs:command*
+    end)
+
+private def mkToJsonInstance (declName : Name) : TermElabM (Array Command) := do
+  let ctx ← mkContext "toJson" declName
+  let cmds := #[← mkToJsonMutualBlock ctx] ++ (← mkInstanceCmds ctx ``ToJson #[declName])
+  trace[Elab.Deriving.toJson] "\n{cmds}"
+  return cmds
+
+private def mkFromJsonInstance (declName : Name) : TermElabM (Array Command) := do
+  let ctx ← mkContext "fromJson" declName
+  let cmds := #[← mkFromJsonMutualBlock ctx] ++ (← mkInstanceCmds ctx ``FromJson #[declName])
+  trace[Elab.Deriving.fromJson] "\n{cmds}"
+  return cmds
+
+def mkToJsonInstanceHandler (declNames : Array Name) : CommandElabM Bool := do
+  if (← declNames.allM isInductive) && declNames.size > 0 then
+    for declName in declNames do
+      let cmds ← liftTermElabM <| mkToJsonInstance declName
+      cmds.forM elabCommand
+    return true
+  else
+    return false
+
+def mkFromJsonInstanceHandler (declNames : Array Name) : CommandElabM Bool := do
+  if (← declNames.allM isInductive) && declNames.size > 0 then
+    for declName in declNames do
+      let cmds ← liftTermElabM <| mkFromJsonInstance declName
+      cmds.forM elabCommand
+    return true
+  else
+    return false
+
+builtin_initialize
+  registerDerivingHandler ``ToJson mkToJsonInstanceHandler
+  registerDerivingHandler ``FromJson mkFromJsonInstanceHandler
+
+  registerTraceClass `Elab.Deriving.toJson
+  registerTraceClass `Elab.Deriving.fromJson
+
+end Lean.Elab.Deriving.FromToJson

backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Elab/Term.lean

@@ -0,0 +1,2128 @@
+/-
+Copyright (c) 2019 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura, Sebastian Ullrich
+-/
+prelude
+import Lean.ReservedNameAction
+import Lean.Meta.AppBuilder
+import Lean.Meta.CollectMVars
+import Lean.Meta.Coe
+import Lean.Util.CollectLevelMVars
+import Lean.Linter.Deprecated
+import Lean.Elab.Config
+import Lean.Elab.Level
+import Lean.Elab.DeclModifiers
+import Lean.Elab.PreDefinition.TerminationHint
+import Lean.Elab.DeclarationRange
+import Lean.Elab.WhereFinally
+import Lean.Language.Basic
+import Lean.Elab.InfoTree.InlayHints
+
+namespace Lean.Elab
+
+namespace Term
+
+/-- Saved context for postponed terms and tactics to be executed. -/
+structure SavedContext where
+  declName?  : Option Name
+  options    : Options
+  openDecls  : List OpenDecl
+  macroStack : MacroStack
+  errToSorry : Bool
+  levelNames : List Name
+
+/-- The kind of a tactic metavariable, used for additional error reporting. -/
+inductive TacticMVarKind
+  /-- Standard tactic metavariable, arising from `by ...` syntax. -/
+  | term
+  /-- Tactic metavariable arising from an autoparam for a function application. -/
+  | autoParam (argName : Name)
+  /-- Tactic metavariable arising from an autoparam for a structure field. -/
+  | fieldAutoParam (fieldName structName : Name)
+
+/-- We use synthetic metavariables as placeholders for pending elaboration steps. -/
+inductive SyntheticMVarKind where
+  /--
+  Use typeclass resolution to synthesize value for metavariable.
+  If `extraErrorMsg?` is `some msg`, `msg` contains additional information to include in error messages
+  regarding type class synthesis failure.
+  -/
+  | typeClass (extraErrorMsg? : Option MessageData)
+  /--
+  Use coercion to synthesize value for the metavariable.
+  If synthesis fails, then throws an error.
+  - If `mkErrorMsg?` is provided, then the error `mkErrorMsg expectedType e` is thrown.
+    The `mkErrorMsg` function is allowed to throw an error itself.
+  - Otherwise, throws a default type mismatch error message.
+    If `header?` is not provided, the default header is "type mismatch".
+    If `f?` is provided, then throws an application type mismatch error.
+  -/
+  | coe (header? : Option String) (expectedType : Expr) (e : Expr) (f? : Option Expr)
+      (mkErrorMsg? : Option (MVarId → Expr → Expr → MetaM MessageData))
+  /--
+  Use tactic to synthesize value for metavariable.
+
+  If `delayOnMVars` is true, the tactic will not be executed until the goal is free of unassigned
+  expr metavariables.
+  -/
+  | tactic (tacticCode : Syntax) (ctx : SavedContext) (kind : TacticMVarKind) (delayOnMVars := false)
+  /-- Metavariable represents a hole whose elaboration has been postponed. -/
+  | postponed (ctx : SavedContext)
+  deriving Inhabited
+
+/--
+Convert an "extra" optional error message into a message `"\n{msg}"` (if `some msg`) and `MessageData.nil` (if `none`)
+-/
+def extraMsgToMsg (extraErrorMsg? : Option MessageData) : MessageData :=
+  if let some msg := extraErrorMsg? then m!"\n{msg}" else .nil
+
+instance : ToString SyntheticMVarKind where
+  toString
+    | .typeClass .. => "typeclass"
+    | .coe ..       => "coe"
+    | .tactic ..    => "tactic"
+    | .postponed .. => "postponed"
+
+structure SyntheticMVarDecl where
+  stx : Syntax
+  kind : SyntheticMVarKind
+  deriving Inhabited
+
+/--
+  We can optionally associate an error context with a metavariable (see `MVarErrorInfo`).
+  We have three different kinds of error context.
+-/
+inductive MVarErrorKind where
+  /-- Metavariable for implicit arguments. `ctx` is the parent application,
+  `lctx` is a local context where it is valid (necessary for eta feature for named arguments). -/
+  | implicitArg (lctx : LocalContext) (ctx : Expr)
+  /-- Metavariable for explicit holes provided by the user (e.g., `_` and `?m`) -/
+  | hole
+  /-- "Custom", `msgData` stores the additional error messages. -/
+  | custom (msgData : MessageData)
+  deriving Inhabited
+
+instance : ToString MVarErrorKind where
+  toString
+    | .implicitArg _ _ => "implicitArg"
+    | .hole            => "hole"
+    | .custom _        => "custom"
+
+/--
+  We can optionally associate an error context with metavariables.
+-/
+structure MVarErrorInfo where
+  mvarId    : MVarId
+  ref       : Syntax
+  kind      : MVarErrorKind
+  deriving Inhabited
+
+/--
+When reporting unexpected universe level metavariables, it is useful to localize the errors
+to particular terms, especially at `let` bindings and function binders,
+where universe polymorphism is not permitted.
+-/
+structure LevelMVarErrorInfo where
+  lctx      : LocalContext
+  expr      : Expr
+  ref       : Syntax
+  msgData?  : Option MessageData := none
+  deriving Inhabited
+
+/--
+  Nested `let rec` expressions are eagerly lifted by the elaborator.
+  We store the information necessary for performing the lifting here.
+-/
+structure LetRecToLift where
+  ref            : Syntax
+  fvarId         : FVarId
+  attrs          : Array Attribute
+  shortDeclName  : Name
+  declName       : Name
+  parentName?    : Option Name
+  lctx           : LocalContext
+  localInstances : LocalInstances
+  type           : Expr
+  val            : Expr
+  mvarId         : MVarId
+  termination    : TerminationHints
+  deriving Inhabited
+
+/--
+  State of the `TermElabM` monad.
+-/
+structure State where
+  levelNames        : List Name       := []
+  syntheticMVars    : MVarIdMap SyntheticMVarDecl := {}
+  pendingMVars      : List MVarId := {}
+  /-- List of errors associated to a metavariable that are shown to the user if the metavariable could not be fully instantiated -/
+  mvarErrorInfos    : List MVarErrorInfo := []
+  /-- List of data to be able to localize universe level metavariable errors to particular expressions. -/
+  levelMVarErrorInfos   : List LevelMVarErrorInfo := []
+  /--
+    `mvarArgNames` stores the argument names associated to metavariables.
+    These are used in combination with `mvarErrorInfos` for throwing errors about metavariables that could not be fully instantiated.
+    For example when elaborating `List _`, the argument name of the placeholder will be `α`.
+
+    While elaborating an application, `mvarArgNames` is set for each metavariable argument, using the available argument name.
+    This may happen before or after the `mvarErrorInfos` is set for the same metavariable.
+
+    We used to store the argument names in `mvarErrorInfos`, updating the `MVarErrorInfos` to add the argument name when it is available,
+    but this doesn't work if the argument name is available _before_ the `mvarErrorInfos` is set for that metavariable.
+  -/
+  mvarArgNames      : MVarIdMap Name := {}
+  letRecsToLift     : List LetRecToLift := []
+  deriving Inhabited
+
+/--
+  Backtrackable state for the `TermElabM` monad.
+-/
+structure SavedState where
+  «meta» : Meta.SavedState
+  «elab» : State
+  deriving Nonempty
+
+end Term
+
+namespace Tactic
+
+/--
+  State of the `TacticM` monad.
+-/
+structure State where
+  goals : List MVarId
+  deriving Inhabited
+
+/--
+  Snapshots are used to implement the `save` tactic.
+  This tactic caches the state of the system, and allows us to "replay"
+  expensive proofs efficiently. This is only relevant implementing the
+  LSP server.
+-/
+structure Snapshot where
+  core   : Core.State
+  «meta» : Meta.State
+  term   : Term.State
+  tactic : Tactic.State
+  stx    : Syntax
+
+/--
+  Key for the cache used to implement the `save` tactic.
+-/
+structure CacheKey where
+  mvarId : MVarId -- TODO: should include all goals
+  pos    : String.Pos
+  deriving BEq, Hashable, Inhabited
+
+/--
+  Cache for the `save` tactic.
+-/
+structure Cache where
+   pre  : PHashMap CacheKey Snapshot := {}
+   post : PHashMap CacheKey Snapshot := {}
+   deriving Inhabited
+
+section Snapshot
+open Language
+
+structure SavedState where
+  term   : Term.SavedState
+  tactic : State
+
+/-- Snapshot after finishing execution of a tactic. -/
+structure TacticFinishedSnapshot extends Language.Snapshot where
+  /-- State saved for reuse, if no fatal exception occurred. -/
+  state? : Option SavedState
+  /-- Untyped snapshots from `logSnapshotTask`, saved at this level for cancellation. -/
+  moreSnaps : Array (SnapshotTask SnapshotTree)
+deriving Inhabited
+instance : ToSnapshotTree TacticFinishedSnapshot where
+  toSnapshotTree s := ⟨s.toSnapshot, s.moreSnaps⟩
+
+/-- Snapshot just before execution of a tactic. -/
+structure TacticParsedSnapshot extends Language.Snapshot where
+  /-- Syntax tree of the tactic, stored and compared for incremental reuse. -/
+  stx      : Syntax
+  /-- Task for nested incrementality, if enabled for tactic. -/
+  inner?   : Option (SnapshotTask TacticParsedSnapshot) := none
+  /-- Task for state after tactic execution. -/
+  finished : SnapshotTask TacticFinishedSnapshot
+  /-- Tasks for subsequent, potentially parallel, tactic steps. -/
+  next     : Array (SnapshotTask TacticParsedSnapshot) := #[]
+deriving Inhabited
+partial instance : ToSnapshotTree TacticParsedSnapshot where
+  toSnapshotTree := go where
+    go := fun s => ⟨s.toSnapshot,
+      s.inner?.toArray.map (·.map (sync := true) go) ++
+      #[s.finished.map (sync := true) toSnapshotTree] ++
+      s.next.map (·.map (sync := true) go)⟩
+
+end Snapshot
+end Tactic
+
+namespace Term
+
+structure Context where
+  declName? : Option Name := none
+  macroStack        : MacroStack      := []
+  /--
+     When `mayPostpone == true`, an elaboration function may interrupt its execution by throwing `Exception.postpone`.
+     The function `elabTerm` catches this exception and creates fresh synthetic metavariable `?m`, stores `?m` in
+     the list of pending synthetic metavariables, and returns `?m`. -/
+  mayPostpone : Bool := true
+  /--
+     When `errToSorry` is set to true, the method `elabTerm` catches
+     exceptions and converts them into synthetic `sorry`s.
+     The implementation of choice nodes and overloaded symbols rely on the fact
+     that when `errToSorry` is set to false for an elaboration function `F`, then
+     `errToSorry` remains `false` for all elaboration functions invoked by `F`.
+     That is, it is safe to transition `errToSorry` from `true` to `false`, but
+     we must not set `errToSorry` to `true` when it is currently set to `false`. -/
+  errToSorry : Bool := true
+  /--
+     When `autoBoundImplicit` is set to true, instead of producing
+     an "unknown identifier" error for unbound variables, we generate an
+     internal exception. This exception is caught at `elabBinders` and
+     `elabTypeWithUnboldImplicit`. Both methods add implicit declarations
+     for the unbound variable and try again. -/
+  autoBoundImplicit  : Bool            := false
+  autoBoundImplicits : PArray Expr := {}
+  /--
+    A name `n` is only eligible to be an auto implicit name if `autoBoundImplicitForbidden n = false`.
+    We use this predicate to disallow `f` to be considered an auto implicit name in a definition such
+    as
+    ```
+    def f : f → Bool := fun _ => true
+    ```
+  -/
+  autoBoundImplicitForbidden : Name → Bool := fun _ => false
+  /-- Map from user name to internal unique name -/
+  sectionVars        : NameMap Name    := {}
+  /-- Map from internal name to fvar -/
+  sectionFVars       : NameMap Expr    := {}
+  /-- Enable/disable implicit lambdas feature. -/
+  implicitLambda     : Bool            := true
+  /-- Heed `elab_as_elim` attribute. -/
+  heedElabAsElim     : Bool            := true
+  /-- Noncomputable sections automatically add the `noncomputable` modifier to any declaration we cannot generate code for. -/
+  isNoncomputableSection : Bool        := false
+  /-- When `true` we skip TC failures. We use this option when processing patterns. -/
+  ignoreTCFailures : Bool := false
+  /-- `true` when elaborating patterns. It affects how we elaborate named holes. -/
+  inPattern        : Bool := false
+  /--
+  Snapshot for incremental processing of current tactic, if any.
+
+  Invariant: if the bundle's `old?` is set, then the state *up to the start* of the tactic is
+  unchanged, i.e. reuse is possible.
+  -/
+  tacSnap?         : Option (Language.SnapshotBundle Tactic.TacticParsedSnapshot) := none
+  /--
+  If `true`, we store in the `Expr` the `Syntax` for recursive applications (i.e., applications
+  of free variables tagged with `isAuxDecl`). We store the `Syntax` using `mkRecAppWithSyntax`.
+  We use the `Syntax` object to produce better error messages at `Structural.lean` and `WF.lean`. -/
+  saveRecAppSyntax : Bool := true
+  /--
+  If `holesAsSyntheticOpaque` is `true`, then we mark metavariables associated
+  with `_`s as `syntheticOpaque` if they do not occur in patterns.
+  This option is useful when elaborating terms in tactics such as `refine'` where
+  we want holes there to become new goals. See issue #1681, we have
+  `refine' (fun x => _)
+  -/
+  holesAsSyntheticOpaque : Bool := false
+  /--
+  If `checkDeprecated := true`, then `Linter.checkDeprecated` when creating constants.
+  -/
+  checkDeprecated : Bool := true
+
+abbrev TermElabM := ReaderT Context $ StateRefT State MetaM
+abbrev TermElab  := Syntax → Option Expr → TermElabM Expr
+
+/-
+Make the compiler generate specialized `pure`/`bind` so we do not have to optimize through the
+whole monad stack at every use site. May eventually be covered by `deriving`.
+-/
+@[always_inline]
+instance : Monad TermElabM :=
+  let i := inferInstanceAs (Monad TermElabM)
+  { pure := i.pure, bind := i.bind }
+
+open Meta
+
+instance : Inhabited (TermElabM α) where
+  default := throw default
+
+protected def saveState : TermElabM SavedState :=
+  return { «meta» := (← Meta.saveState), «elab» := (← get) }
+
+def SavedState.restore (s : SavedState) (restoreInfo : Bool := false) : TermElabM Unit := do
+  let traceState ← getTraceState -- We never backtrack trace message
+  let infoState ← getInfoState -- We also do not backtrack the info nodes when `restoreInfo == false`
+  s.meta.restore
+  set s.elab
+  setTraceState traceState
+  unless restoreInfo do
+    setInfoState infoState
+
+/--
+Like `Meta.withRestoreOrSaveFull` for `TermElabM`, but also takes a `tacSnap?` that
+* when running `act`, is set as `Context.tacSnap?`
+* otherwise (i.e. on restore) is used to update the new snapshot promise to the old task's
+  value.
+This extra restore step is necessary because while `reusableResult?` can be used to replay any
+effects on `State`, `Context.tacSnap?` is not part of it but changed via an `IO` side effect, so
+it needs to be replayed separately.
+
+We use an explicit parameter instead of accessing `Context.tacSnap?` directly because this prevents
+`withRestoreOrSaveFull` and `withReader` from being used in the wrong order.
+-/
+@[specialize]
+def withRestoreOrSaveFull (reusableResult? : Option (α × SavedState))
+    (tacSnap? : Option (Language.SnapshotBundle Tactic.TacticParsedSnapshot)) (act : TermElabM α) :
+    TermElabM (α × SavedState) := do
+  if let some (_, state) := reusableResult? then
+    set state.elab
+    if let some snap := tacSnap? then
+      let some old := snap.old?
+        | throwError "withRestoreOrSaveFull: expected old snapshot in `tacSnap?`"
+      snap.new.resolve old.val.get
+
+  let reusableResult? := reusableResult?.map (fun (val, state) => (val, state.meta))
+  let (a, «meta») ← withReader ({ · with tacSnap? }) do
+    controlAt MetaM fun runInBase => do
+      Meta.withRestoreOrSaveFull reusableResult? <| runInBase act
+  return (a, { «meta», «elab» := (← get) })
+
+instance : MonadBacktrack SavedState TermElabM where
+  saveState      := Term.saveState
+  restoreState b := b.restore
+
+/--
+Incremental elaboration helper. Avoids leakage of data from outside syntax via the monadic context
+when running `act` on `stx` by
+* setting `stx` as the `ref` and
+* deactivating `suppressElabErrors` if `stx` is `missing`-free, which also helps with not hiding
+  useful errors in this part of the input. Note that if `stx` has `missing`, this should always be
+  true for the outer syntax as well, so taking the old value of `suppressElabErrors` into account
+  should not introduce data leakage.
+
+This combinator should always be used when narrowing reuse to a syntax subtree, usually (in the case
+of tactics, to be generalized) via `withNarrowed(Arg)TacticReuse`.
+-/
+def withReuseContext [Monad m] [MonadWithReaderOf Core.Context m] (stx : Syntax) (act : m α) :
+    m α := do
+  withTheReader Core.Context (fun ctx => { ctx with
+      ref := stx
+      suppressElabErrors := ctx.suppressElabErrors && stx.hasMissing }) act
+
+/--
+Manages reuse information for nested tactics by `split`ting given syntax into an outer and inner
+part. `act` is then run on the inner part but with reuse information adjusted as following:
+* If the old (from `tacSnap?`'s `SyntaxGuarded.stx`) and new (from `stx`) outer syntax are not
+  identical according to `Syntax.eqWithInfo`, reuse is disabled.
+* Otherwise, the old syntax as stored in `tacSnap?` is updated to the old *inner* syntax.
+* In any case, `withReuseContext` is used on the new inner syntax to further prepare the monadic
+  context.
+
+For any tactic that participates in reuse, `withNarrowedTacticReuse` should be applied to the
+tactic's syntax and `act` should be used to do recursive tactic evaluation of nested parts. Also,
+after this function, `getAndEmptySnapshotTasks` should be called and the result stored in a snapshot
+so that the tasks don't end up in a snapshot further up and are cancelled together with it; see
+note [Incremental Cancellation].
+-/
+def withNarrowedTacticReuse [Monad m] [MonadReaderOf Context m] [MonadLiftT BaseIO m]
+    [MonadWithReaderOf Core.Context m] [MonadWithReaderOf Context m] [MonadOptions m]
+    (split : Syntax → Syntax × Syntax) (act : Syntax → m α) (stx : Syntax) : m α := do
+  let (outer, inner) := split stx
+  let opts ← getOptions
+  let ctx ← readThe Term.Context
+  withTheReader Term.Context (fun ctx => { ctx with tacSnap? := ctx.tacSnap?.map fun tacSnap =>
+    { tacSnap with old? := tacSnap.old?.bind fun old => do
+      let (oldOuter, oldInner) := split old.stx
+      guard <| outer.eqWithInfoAndTraceReuse opts oldOuter
+      return { old with stx := oldInner }
+    }
+  }) do
+    if let some oldOuter := ctx.tacSnap?.bind (·.old?) then
+      if (← read).tacSnap?.bind (·.old?) |>.isNone then
+        oldOuter.val.cancelRec
+    withReuseContext inner (act inner)
+
+/--
+A variant of `withNarrowedTacticReuse` that uses `stx[argIdx]` as the inner syntax and all `stx`
+child nodes before that as the outer syntax, i.e. reuse is disabled if there was any change before
+`argIdx`.
+
+NOTE: child nodes after `argIdx` are not tested (which would almost always disable reuse as they are
+necessarily shifted by changes at `argIdx`) so it must be ensured that the result of `arg` does not
+depend on them (i.e. they should not be inspected beforehand).
+-/
+def withNarrowedArgTacticReuse [Monad m] [MonadReaderOf Context m] [MonadLiftT BaseIO m]
+    [MonadWithReaderOf Core.Context m] [MonadWithReaderOf Context m] [MonadOptions m]
+    (argIdx : Nat) (act : Syntax → m α) (stx : Syntax) : m α :=
+  withNarrowedTacticReuse (fun stx => (mkNullNode stx.getArgs[*...argIdx], stx[argIdx])) act stx
+
+/--
+Disables incremental tactic reuse *and* reporting for `act` if `cond` is true by setting `tacSnap?`
+to `none`. This should be done for tactic blocks that are run multiple times as otherwise the
+reported progress will jump back and forth (and partial reuse for these kinds of tact blocks is
+similarly questionable).
+-/
+def withoutTacticIncrementality [Monad m] [MonadWithReaderOf Context m] [MonadOptions m]
+    (cond : Bool) (act : m α) : m α := do
+  let opts ← getOptions
+  withTheReader Term.Context (fun ctx => { ctx with tacSnap? := ctx.tacSnap?.filter fun tacSnap => Id.run do
+    if let some old := tacSnap.old? then
+      if cond && opts.getBool `trace.Elab.reuse then
+        dbg_trace "reuse stopped: guard failed at {old.stx}"
+    return !cond
+  }) act
+
+/-- Disables incremental tactic reuse for `act` if `cond` is true. -/
+def withoutTacticReuse [Monad m] [MonadWithReaderOf Context m] [MonadOptions m]
+    (cond : Bool) (act : m ��) : m α := do
+  let opts ← getOptions
+  withTheReader Term.Context (fun ctx => { ctx with tacSnap? := ctx.tacSnap?.map fun tacSnap =>
+    { tacSnap with old? := tacSnap.old?.filter fun old => Id.run do
+      if cond && opts.getBool `trace.Elab.reuse then
+        dbg_trace "reuse stopped: guard failed at {old.stx}"
+      return !cond }
+  }) act
+
+@[inherit_doc Core.wrapAsyncAsSnapshot]
+def wrapAsyncAsSnapshot {α : Type} (act : α → TermElabM Unit) (cancelTk? : Option IO.CancelToken)
+    (desc : String := by exact decl_name%.toString) :
+    TermElabM (α → BaseIO Language.SnapshotTree) := do
+  let ctx ← read
+  let st ← get
+  let metaCtx ← readThe Meta.Context
+  let metaSt ← getThe Meta.State
+  Core.wrapAsyncAsSnapshot (cancelTk? := cancelTk?) (desc := desc) fun a =>
+    act a |>.run ctx |>.run' st |>.run' metaCtx metaSt
+
+abbrev TermElabResult (α : Type) := EStateM.Result Exception SavedState α
+
+/--
+  Execute `x`, save resulting expression and new state.
+  We remove any `Info` created by `x`.
+  The info nodes are committed when we execute `applyResult`.
+  We use `observing` to implement overloaded notation and decls.
+  We want to save `Info` nodes for the chosen alternative.
+-/
+def observing (x : TermElabM α) : TermElabM (TermElabResult α) := do
+  let s ← saveState
+  try
+    let e ← x
+    let sNew ← saveState
+    s.restore (restoreInfo := true)
+    return EStateM.Result.ok e sNew
+  catch
+    | ex@(.error ..) =>
+      let sNew ← saveState
+      s.restore (restoreInfo := true)
+      return .error ex sNew
+    | ex@(.internal id _) =>
+      if id == postponeExceptionId then
+        s.restore (restoreInfo := true)
+      throw ex
+
+/--
+  Apply the result/exception and state captured with `observing`.
+  We use this method to implement overloaded notation and symbols. -/
+def applyResult (result : TermElabResult α) : TermElabM α := do
+  match result with
+  | .ok a r     => r.restore (restoreInfo := true); return a
+  | .error ex r => r.restore (restoreInfo := true); throw ex
+
+/--
+  Execute `x`, but keep state modifications only if `x` did not postpone.
+  This method is useful to implement elaboration functions that cannot decide whether
+  they need to postpone or not without updating the state. -/
+def commitIfDidNotPostpone (x : TermElabM α) : TermElabM α := do
+  -- We just reuse the implementation of `observing` and `applyResult`.
+  let r ← observing x
+  applyResult r
+
+/--
+  Return the universe level names explicitly provided by the user.
+-/
+def getLevelNames : TermElabM (List Name) :=
+  return (← get).levelNames
+
+/--
+  Given a free variable `fvar`, return its declaration.
+  This function panics if `fvar` is not a free variable.
+-/
+def getFVarLocalDecl! (fvar : Expr) : TermElabM LocalDecl := do
+  match (← getLCtx).find? fvar.fvarId! with
+  | some d => pure d
+  | none   => unreachable!
+
+instance : AddErrorMessageContext TermElabM where
+  add ref msg := do
+    let ctx ← read
+    let ref := getBetterRef ref ctx.macroStack
+    let msg ← addMessageContext msg
+    let msg ← addMacroStack msg ctx.macroStack
+    pure (ref, msg)
+
+/--
+  Execute `x` without storing `Syntax` for recursive applications. See `saveRecAppSyntax` field at `Context`.
+-/
+def withoutSavingRecAppSyntax (x : TermElabM α) : TermElabM α :=
+  withReader (fun ctx => { ctx with saveRecAppSyntax := false }) x
+
+unsafe def mkTermElabAttributeUnsafe (ref : Name) : IO (KeyedDeclsAttribute TermElab) :=
+  mkElabAttribute TermElab `builtin_term_elab `term_elab `Lean.Parser.Term `Lean.Elab.Term.TermElab "term" ref
+
+@[implemented_by mkTermElabAttributeUnsafe]
+opaque mkTermElabAttribute (ref : Name) : IO (KeyedDeclsAttribute TermElab)
+
+/--
+Registers a term elaborator for the given syntax node kind.
+
+A term elaborator should have type `Lean.Elab.Term.TermElab` (which is
+`Lean.Syntax → Option Lean.Expr → Lean.Elab.Term.TermElabM Lean.Expr`), i.e. should take syntax of
+the given syntax node kind and an optional expected type as parameters and produce an expression.
+
+The `elab_rules` and `elab` commands should usually be preferred over using this attribute
+directly.
+-/
+@[builtin_doc]
+builtin_initialize termElabAttribute : KeyedDeclsAttribute TermElab ← mkTermElabAttribute decl_name%
+
+/--
+  Auxiliary datatype for presenting a Lean lvalue modifier.
+  We represent an unelaborated lvalue as a `Syntax` (or `Expr`) and `List LVal`.
+  Example: `a.foo.1` is represented as the `Syntax` `a` and the list
+  `[LVal.fieldName "foo", LVal.fieldIdx 1]`.
+-/
+inductive LVal where
+  | fieldIdx  (ref : Syntax) (i : Nat)
+  /-- Field `suffix?` is for producing better error messages because `x.y` may be a field access or a hierarchical/composite name.
+  `ref` is the syntax object representing the field. `fullRef` includes the LHS. -/
+  | fieldName (ref : Syntax) (name : String) (suffix? : Option Name) (fullRef : Syntax)
+
+def LVal.getRef : LVal → Syntax
+  | .fieldIdx ref _    => ref
+  | .fieldName ref ..  => ref
+
+def LVal.isFieldName : LVal ��� Bool
+  | .fieldName .. => true
+  | _ => false
+
+instance : ToString LVal where
+  toString
+    | .fieldIdx _ i     => toString i
+    | .fieldName _ n .. => n
+
+/-- Return the name of the declaration being elaborated if available. -/
+def getDeclName? : TermElabM (Option Name) := return (← read).declName?
+/-- Return the list of nested `let rec` declarations that need to be lifted. -/
+def getLetRecsToLift : TermElabM (List LetRecToLift) := return (← get).letRecsToLift
+/-- Return the declaration of the given metavariable -/
+def getMVarDecl (mvarId : MVarId) : TermElabM MetavarDecl := return (← getMCtx).getDecl mvarId
+
+instance : MonadParentDecl TermElabM where
+  getParentDeclName? := getDeclName?
+
+/--
+Executes `x` in the context of the given declaration name. Ensures that the info tree is set up
+correctly and adjusts the declaration name generator to generate names below this name, resetting
+the nested counter.
+-/
+def withDeclName (name : Name) (x : TermElabM α) : TermElabM α :=
+  withReader (fun ctx => { ctx with declName? := name }) do
+  withSaveParentDeclInfoContext do
+  withDeclNameForAuxNaming name do
+    x
+
+/-- Update the universe level parameter names. -/
+def setLevelNames (levelNames : List Name) : TermElabM Unit :=
+  modify fun s => { s with levelNames := levelNames }
+
+/-- Execute `x` using `levelNames` as the universe level parameter names. See `getLevelNames`. -/
+def withLevelNames (levelNames : List Name) (x : TermElabM α) : TermElabM α := do
+  let levelNamesSaved ← getLevelNames
+  setLevelNames levelNames
+  try x finally setLevelNames levelNamesSaved
+
+def withoutErrToSorryImp (x : TermElabM α) : TermElabM α :=
+  withReader (fun ctx => { ctx with errToSorry := false }) x
+
+/--
+  Execute `x` without converting errors (i.e., exceptions) to `sorry` applications.
+  Recall that when `errToSorry = true`, the method `elabTerm` catches exceptions and converts them into `sorry` applications.
+-/
+def withoutErrToSorry [MonadFunctorT TermElabM m] : m α → m α :=
+  monadMap (m := TermElabM) withoutErrToSorryImp
+
+def withoutHeedElabAsElimImp (x : TermElabM α) : TermElabM α :=
+  withReader (fun ctx => { ctx with heedElabAsElim := false }) x
+
+/--
+  Execute `x` without heeding the `elab_as_elim` attribute. Useful when there is
+  no expected type (so `elabAppArgs` would fail), but expect that the user wants
+  to use such constants.
+-/
+def withoutHeedElabAsElim [MonadFunctorT TermElabM m] : m α → m α :=
+  monadMap (m := TermElabM) withoutHeedElabAsElimImp
+
+/--
+Execute `x` but discard changes performed at `Term.State` and `Meta.State`.
+Recall that the `Environment`, `InfoState` and messages are at `Core.State`.  Thus, any updates to
+it will be preserved.
+This method is useful for performing computations where all metavariable must be resolved or
+discarded.
+The `InfoTree`s are not discarded, however, and wrapped in `InfoTree.Context`
+to store their metavariable context.
+-/
+def withoutModifyingElabMetaStateWithInfo (x : TermElabM α) : TermElabM α := do
+  let s ← get
+  let sMeta ← getThe Meta.State
+  try
+    withSaveInfoContext x
+  finally
+    set s
+    set sMeta
+
+/--
+  Execute `x` but discard changes performed to the state.
+  However, the info trees and messages are not discarded. -/
+private def withoutModifyingStateWithInfoAndMessagesImpl (x : TermElabM α) : TermElabM α := do
+  let saved ← saveState
+  try
+    withSaveInfoContext x
+  finally
+    let saved := { saved with meta.core.infoState := (← getInfoState), meta.core.messages := (← getThe Core.State).messages }
+    restoreState saved
+
+/--
+Wraps the trees returned from `getInfoTrees`, if any, in an `InfoTree.context` node based on the
+current monadic context and state. This is mainly used to report info trees early via
+`Snapshot.infoTree?`. The trees are not removed from the `getInfoTrees` state as the final info tree
+of the elaborated command should be complete and not depend on whether parts have been reported
+early.
+
+As `InfoTree.context` can have only one child, this function panics if `trees` contains more than 1
+tree. Also, `PartialContextInfo.parentDeclCtx` is not currently generated as that information is not
+available in the monadic context and only needed for the final info tree.
+-/
+def getInfoTreeWithContext? : TermElabM (Option InfoTree) := do
+  let st ← getInfoState
+  if st.trees.size > 1 then
+    return panic! "getInfoTreeWithContext: overfull tree"
+  let some t := st.trees[0]? |
+    return none
+  let t := t.substitute st.assignment
+  let ctx ← readThe Core.Context
+  let s ← getThe Core.State
+  let ctx := PartialContextInfo.commandCtx {
+    env := s.env, fileMap := ctx.fileMap, mctx := {}, currNamespace := ctx.currNamespace,
+    openDecls := ctx.openDecls, options := ctx.options, ngen := s.ngen
+  }
+  return InfoTree.context ctx t
+
+/-- For testing `TermElabM` methods. The #eval command will sign the error. -/
+def throwErrorIfErrors : TermElabM Unit := do
+  if (← MonadLog.hasErrors) then
+    throwError "Error(s)"
+
+def traceAtCmdPos (cls : Name) (msg : Unit → MessageData) : TermElabM Unit :=
+  withRef Syntax.missing <| trace cls msg
+
+def ppGoal (mvarId : MVarId) : TermElabM Format :=
+  Meta.ppGoal mvarId
+
+open Level (LevelElabM)
+
+def liftLevelM (x : LevelElabM α) : TermElabM α := do
+  let ctx ← read
+  let mctx ← getMCtx
+  let ngen ← getNGen
+  let lvlCtx : Level.Context := { options := (← getOptions), ref := (← getRef), autoBoundImplicit := ctx.autoBoundImplicit }
+  match (x lvlCtx).run { ngen := ngen, mctx := mctx, levelNames := (← getLevelNames) } with
+  | .ok a newS  => setMCtx newS.mctx; setNGen newS.ngen; setLevelNames newS.levelNames; pure a
+  | .error ex _ => throw ex
+
+def elabLevel (stx : Syntax) : TermElabM Level :=
+  liftLevelM <| Level.elabLevel stx
+
+/-- Elaborate `x` with `stx` on the macro stack -/
+def withPushMacroExpansionStack (beforeStx afterStx : Syntax) (x : TermElabM α) : TermElabM α :=
+  withReader (fun ctx => { ctx with macroStack := { before := beforeStx, after := afterStx } :: ctx.macroStack }) x
+
+/-- Elaborate `x` with `stx` on the macro stack and produce macro expansion info -/
+def withMacroExpansion (beforeStx afterStx : Syntax) (x : TermElabM α) : TermElabM α :=
+  withMacroExpansionInfo beforeStx afterStx do
+    withPushMacroExpansionStack beforeStx afterStx x
+
+/--
+  Add the given metavariable to the list of pending synthetic metavariables.
+  The method `synthesizeSyntheticMVars` is used to process the metavariables on this list. -/
+def registerSyntheticMVar (stx : Syntax) (mvarId : MVarId) (kind : SyntheticMVarKind) : TermElabM Unit := do
+  modify fun s => { s with syntheticMVars := s.syntheticMVars.insert mvarId { stx, kind }, pendingMVars := mvarId :: s.pendingMVars }
+
+def registerSyntheticMVarWithCurrRef (mvarId : MVarId) (kind : SyntheticMVarKind) : TermElabM Unit := do
+  registerSyntheticMVar (← getRef) mvarId kind
+
+def registerMVarErrorInfo (mvarErrorInfo : MVarErrorInfo) : TermElabM Unit :=
+  modify fun s => { s with mvarErrorInfos := mvarErrorInfo :: s.mvarErrorInfos }
+
+def registerMVarErrorHoleInfo (mvarId : MVarId) (ref : Syntax) : TermElabM Unit :=
+  registerMVarErrorInfo { mvarId, ref, kind := .hole }
+
+def registerMVarErrorImplicitArgInfo (mvarId : MVarId) (ref : Syntax) (app : Expr) : TermElabM Unit := do
+  registerMVarErrorInfo { mvarId, ref, kind := .implicitArg (← getLCtx) app }
+
+def registerMVarErrorCustomInfo (mvarId : MVarId) (ref : Syntax) (msgData : MessageData) : TermElabM Unit := do
+  registerMVarErrorInfo { mvarId, ref, kind := .custom msgData }
+
+def registerCustomErrorIfMVar (e : Expr) (ref : Syntax) (msgData : MessageData) : TermElabM Unit :=
+  match e.getAppFn with
+  | Expr.mvar mvarId => registerMVarErrorCustomInfo mvarId ref msgData
+  | _ => pure ()
+
+def registerMVarArgName (mvarId : MVarId) (argName : Name) : TermElabM Unit :=
+  modify fun s => { s with mvarArgNames := s.mvarArgNames.insert mvarId argName }
+
+/--
+  Auxiliary method for reporting errors of the form "... contains metavariables ...".
+  This kind of error is thrown, for example, at `Match.lean` where elaboration
+  cannot continue if there are metavariables in patterns.
+  We only want to log it if we haven't logged any errors so far. -/
+def throwMVarError (m : MessageData) : TermElabM α := do
+  if (← MonadLog.hasErrors) then
+    throwAbortTerm
+  else
+    throwError m
+
+def MVarErrorInfo.logError (mvarErrorInfo : MVarErrorInfo) (extraMsg? : Option MessageData) : TermElabM Unit := do
+  match mvarErrorInfo.kind with
+  | MVarErrorKind.implicitArg lctx app => withLCtx lctx {} do
+    let app ← instantiateMVars app
+    let msg ← addArgName "don't know how to synthesize implicit argument"
+    let msg := msg ++ m!"{indentExpr app.setAppPPExplicitForExposingMVars}" ++ Format.line ++ "context:" ++ Format.line ++ MessageData.ofGoal mvarErrorInfo.mvarId
+    logErrorAt mvarErrorInfo.ref (appendExtra msg)
+  | MVarErrorKind.hole => do
+    let msg ← addArgName "don't know how to synthesize placeholder" " for argument"
+    let msg := msg ++ Format.line ++ "context:" ++ Format.line ++ MessageData.ofGoal mvarErrorInfo.mvarId
+    logErrorAt mvarErrorInfo.ref (MessageData.tagged `Elab.synthPlaceholder <| appendExtra msg)
+  | MVarErrorKind.custom msg =>
+    logErrorAt mvarErrorInfo.ref (appendExtra msg)
+where
+  /-- Append the argument name (if available) to the message.
+      Remark: if the argument name contains macro scopes we do not append it. -/
+  addArgName (msg : MessageData) (extra : String := "") : TermElabM MessageData := do
+    match (← get).mvarArgNames.find? mvarErrorInfo.mvarId with
+    | none => return msg
+    | some argName => return if argName.hasMacroScopes then msg else msg ++ extra ++ m!" '{argName}'"
+
+  appendExtra (msg : MessageData) : MessageData :=
+    match extraMsg? with
+    | none => msg
+    | some extraMsg => msg ++ extraMsg
+
+/--
+  Try to log errors for the unassigned metavariables `pendingMVarIds`.
+
+  Return `true` if there were "unfilled holes", and we should "abort" declaration.
+  TODO: try to fill "all" holes using synthetic "sorry's"
+
+  Remark: We only log the "unfilled holes" as new errors if no error has been logged so far. -/
+def logUnassignedUsingErrorInfos (pendingMVarIds : Array MVarId) (extraMsg? : Option MessageData := none) : TermElabM Bool := do
+  if pendingMVarIds.isEmpty then
+    return false
+  else
+    let hasOtherErrors ← MonadLog.hasErrors
+    let mut hasNewErrors := false
+    let mut alreadyVisited : MVarIdSet := {}
+    let mut errors : Array MVarErrorInfo := #[]
+    for mvarErrorInfo in (← get).mvarErrorInfos do
+      let mvarId := mvarErrorInfo.mvarId
+      unless alreadyVisited.contains mvarId do
+        alreadyVisited := alreadyVisited.insert mvarId
+        /- The metavariable `mvarErrorInfo.mvarId` may have been assigned or
+           delayed assigned to another metavariable that is unassigned. -/
+        let mvarDeps ← getMVars (mkMVar mvarId)
+        if mvarDeps.any pendingMVarIds.contains then do
+          unless hasOtherErrors do
+            errors := errors.push mvarErrorInfo
+          hasNewErrors := true
+    -- To sort the errors by position use
+    -- let sortedErrors := errors.qsort fun e₁ e₂ => e₁.ref.getPos?.getD 0 < e₂.ref.getPos?.getD 0
+    for error in errors do
+      error.mvarId.withContext do
+        error.logError extraMsg?
+    return hasNewErrors
+
+def registerLevelMVarErrorInfo (levelMVarErrorInfo : LevelMVarErrorInfo) : TermElabM Unit :=
+  modify fun s => { s with levelMVarErrorInfos := levelMVarErrorInfo :: s.levelMVarErrorInfos }
+
+def registerLevelMVarErrorExprInfo (expr : Expr) (ref : Syntax) (msgData? : Option MessageData := none) : TermElabM Unit := do
+  registerLevelMVarErrorInfo { lctx := (← getLCtx), expr, ref, msgData? }
+
+def exposeLevelMVars (e : Expr) : MetaM Expr :=
+  Core.transform e
+    (post := fun e => do
+      match e with
+      | .const _ us     => return .done <| if us.any (·.isMVar) then e.setPPUniverses true else e
+      | .sort u         => return .done <| if u.isMVar then e.setPPUniverses true else e
+      | .lam _ t _ _    => return .done <| if t.hasLevelMVar then e.setOption `pp.funBinderTypes true else e
+      | .letE _ t _ _ _ => return .done <| if t.hasLevelMVar then e.setOption `pp.letVarTypes true else e
+      | _               => return .done e)
+
+def LevelMVarErrorInfo.logError (levelMVarErrorInfo : LevelMVarErrorInfo) : TermElabM Unit :=
+  Meta.withLCtx levelMVarErrorInfo.lctx {} do
+    let e' ← exposeLevelMVars (← instantiateMVars levelMVarErrorInfo.expr)
+    let msg := levelMVarErrorInfo.msgData?.getD m!"don't know how to synthesize universe level metavariables"
+    let msg := m!"{msg}{indentExpr e'}"
+    logErrorAt levelMVarErrorInfo.ref msg
+
+/--
+Try to log errors for unassigned level metavariables `pendingLevelMVarIds`.
+
+Returns `true` if there are any relevant `LevelMVarErrorInfo`s and we should "abort" the declaration.
+
+Remark: we only log unassigned level metavariables as new errors if no error has been logged so far.
+-/
+def logUnassignedLevelMVarsUsingErrorInfos (pendingLevelMVarIds : Array LMVarId) : TermElabM Bool := do
+  if pendingLevelMVarIds.isEmpty then
+    return false
+  else
+    let hasOtherErrors ← MonadLog.hasErrors
+    let mut hasNewErrors := false
+    let mut errors : Array LevelMVarErrorInfo := #[]
+    for levelMVarErrorInfo in (← get).levelMVarErrorInfos do
+      let e ← instantiateMVars levelMVarErrorInfo.expr
+      let lmvars := (collectLevelMVars {} e).result
+      if lmvars.any pendingLevelMVarIds.contains then do
+        unless hasOtherErrors do
+          errors := errors.push levelMVarErrorInfo
+        hasNewErrors := true
+    for error in errors do
+      error.logError
+    return hasNewErrors
+
+/-- Ensure metavariables registered using `registerMVarErrorInfos` (and used in the given declaration) have been assigned. -/
+def ensureNoUnassignedMVars (decl : Declaration) : TermElabM Unit := do
+  let pendingMVarIds ← getMVarsAtDecl decl
+  if (← logUnassignedUsingErrorInfos pendingMVarIds) then
+    throwAbortCommand
+
+/--
+  Execute `x` without allowing it to postpone elaboration tasks.
+  That is, `tryPostpone` is a noop. -/
+def withoutPostponing (x : TermElabM α) : TermElabM α :=
+  withReader (fun ctx => { ctx with mayPostpone := false }) x
+
+/-- Creates syntax for `(` <ident> `:` <type> `)` -/
+def mkExplicitBinder (ident : Syntax) (type : Syntax) : Syntax :=
+  mkNode ``Lean.Parser.Term.explicitBinder #[mkAtom "(", mkNullNode #[ident], mkNullNode #[mkAtom ":", type], mkNullNode, mkAtom ")"]
+
+/--
+  Convert unassigned universe level metavariables into parameters.
+  The new parameter names are fresh names of the form `u_i` with regard to `ctx.levelNames`, which is updated with the new names. -/
+def levelMVarToParam (e : Expr) (except : LMVarId → Bool := fun _ => false) : TermElabM Expr := do
+  let levelNames ← getLevelNames
+  let r := (← getMCtx).levelMVarToParam (fun n => levelNames.elem n) except e `u 1
+  -- Recall that the most recent universe is the first element of the field `levelNames`.
+  setLevelNames (r.newParamNames.reverse.toList ++ levelNames)
+  setMCtx r.mctx
+  return r.expr
+
+/--
+Creates a fresh inaccessible binder name based on `x`.
+Equivalent to ``Lean.Core.mkFreshUserName `x``.
+
+Do not confuse with `Lean.mkFreshId`, for creating fresh free variable and metavariable ids.
+-/
+def mkFreshBinderName [Monad m] [MonadQuotation m] : m Name :=
+  withFreshMacroScope <| MonadQuotation.addMacroScope `x
+
+/--
+  Auxiliary method for creating a `Syntax.ident` containing
+  a fresh name. This method is intended for creating fresh binder names.
+  It is just a thin layer on top of `mkFreshUserName`. -/
+def mkFreshIdent [Monad m] [MonadQuotation m] (ref : Syntax) (canonical := false) : m Ident :=
+  return mkIdentFrom ref (← mkFreshBinderName) canonical
+
+private def applyAttributesCore
+    (declName : Name) (attrs : Array Attribute)
+    (applicationTime? : Option AttributeApplicationTime) : TermElabM Unit := do profileitM Exception "attribute application" (← getOptions) do
+  /-
+  Remark: if the declaration has syntax errors, `declName` may be `.anonymous` see issue #4309
+  In this case, we skip attribute application.
+  -/
+  if declName == .anonymous then
+    return
+  withDeclName declName do
+    for attr in attrs do
+      withTraceNode `Elab.attribute (fun _ => pure m!"applying [{attr.stx}]") do
+      withRef attr.stx do withLogging do
+      let env ← getEnv
+      match getAttributeImpl env attr.name with
+      | Except.error errMsg => throwError errMsg
+      | Except.ok attrImpl  =>
+        let runAttr := attrImpl.add declName attr.stx attr.kind
+        let runAttr := do
+          -- not truly an elaborator, but a sensible target for go-to-definition
+          let elaborator := attrImpl.ref
+          if (← getInfoState).enabled then
+            withInfoContext (mkInfo := return .ofCommandInfo { elaborator, stx := attr.stx }) do
+              try runAttr
+              finally if attr.stx[0].isIdent || attr.stx[0].isAtom then
+                -- Add an additional node over the leading identifier if there is one to make it look more function-like.
+                -- Do this last because we want user-created infos to take precedence
+                pushInfoLeaf <| .ofCommandInfo { elaborator, stx := attr.stx[0] }
+          else
+            runAttr
+        match applicationTime? with
+        | none => runAttr
+        | some applicationTime =>
+          if applicationTime == attrImpl.applicationTime then
+            runAttr
+
+/-- Apply given attributes **at** a given application time -/
+def applyAttributesAt (declName : Name) (attrs : Array Attribute) (applicationTime : AttributeApplicationTime) : TermElabM Unit :=
+  applyAttributesCore declName attrs applicationTime
+
+def applyAttributes (declName : Name) (attrs : Array Attribute) : TermElabM Unit :=
+  applyAttributesCore declName attrs none
+
+def mkTypeMismatchError (header? : Option MessageData) (e : Expr) (eType : Expr) (expectedType : Expr) : MetaM MessageData := do
+  let header : MessageData := match header? with
+    | some header => m!"{header} "
+    | none        => m!"type mismatch{indentExpr e}\n"
+  return m!"{header}{← mkHasTypeButIsExpectedMsg eType expectedType}"
+
+def throwTypeMismatchError (header? : Option MessageData) (expectedType : Expr) (eType : Expr) (e : Expr)
+    (f? : Option Expr := none) (_extraMsg? : Option MessageData := none) : MetaM α := do
+  /-
+    We ignore `extraMsg?` for now. In all our tests, it contained no useful information. It was
+    always of the form:
+    ```
+    failed to synthesize instance
+      CoeT <eType> <e> <expectedType>
+    ```
+    We should revisit this decision in the future and decide whether it may contain useful information
+    or not. -/
+  let extraMsg := Format.nil
+  /-
+  let extraMsg : MessageData := match extraMsg? with
+    | none          => Format.nil
+    | some extraMsg => Format.line ++ extraMsg;
+  -/
+  match f? with
+  | none   => throwError "{← mkTypeMismatchError header? e eType expectedType}{extraMsg}"
+  | some f => Meta.throwAppTypeMismatch f e
+
+def withoutMacroStackAtErr (x : TermElabM α) : TermElabM α :=
+  withTheReader Core.Context (fun (ctx : Core.Context) => { ctx with options := pp.macroStack.set ctx.options false }) x
+
+namespace ContainsPendingMVar
+
+abbrev M := MonadCacheT Expr Unit (OptionT MetaM)
+
+/-- See `containsPostponedTerm` -/
+partial def visit (e : Expr) : M Unit := do
+  checkCache e fun _ => do
+    match e with
+    | .forallE _ d b _   => visit d; visit b
+    | .lam _ d b _       => visit d; visit b
+    | .letE _ t v b _    => visit t; visit v; visit b
+    | .app f a           => visit f; visit a
+    | .mdata _ b         => visit b
+    | .proj _ _ b        => visit b
+    | .fvar fvarId ..    =>
+      match (← fvarId.getDecl) with
+      | .cdecl .. => return ()
+      | .ldecl (value := v) .. => visit v
+    | .mvar mvarId ..    =>
+      let e' ← instantiateMVars e
+      if e' != e then
+        visit e'
+      else
+        match (← getDelayedMVarAssignment? mvarId) with
+        | some d => visit (mkMVar d.mvarIdPending)
+        | none   => failure
+    | _ => return ()
+
+end ContainsPendingMVar
+
+/-- Return `true` if `e` contains a pending metavariable. Remark: it also visits let-declarations. -/
+def containsPendingMVar (e : Expr) : MetaM Bool := do
+  match (← ContainsPendingMVar.visit e |>.run.run) with
+  | some _ => return false
+  | none   => return true
+
+/--
+  Try to synthesize metavariable using type class resolution.
+  This method assumes the local context and local instances of `instMVar` coincide
+  with the current local context and local instances.
+  Return `true` if the instance was synthesized successfully, and `false` if
+  the instance contains unassigned metavariables that are blocking the type class
+  resolution procedure. Throw an exception if resolution or assignment irrevocably fails.
+
+  If `extraErrorMsg?` is not none, it contains additional information that should be attached
+  to type class synthesis failures.
+-/
+def synthesizeInstMVarCore (instMVar : MVarId) (maxResultSize? : Option Nat := none) (extraErrorMsg? : Option MessageData := none): TermElabM Bool := do
+  let extraErrorMsg := extraMsgToMsg extraErrorMsg?
+  let instMVarDecl ← getMVarDecl instMVar
+  let type := instMVarDecl.type
+  let type ← instantiateMVars type
+  let result ← trySynthInstance type maxResultSize?
+  match result with
+  | LOption.some val =>
+    if (← instMVar.isAssigned) then
+      let oldVal ← instantiateMVars (mkMVar instMVar)
+      unless (← isDefEq oldVal val) do
+        if (← containsPendingMVar oldVal <||> containsPendingMVar val) then
+          /- If `val` or `oldVal` contains metavariables directly or indirectly (e.g., in a let-declaration),
+             we return `false` to indicate we should try again later. This is very coarse grain since
+             the metavariable may not be responsible for the failure. We should refine the test in the future if needed.
+             This check has been added to address dependencies between postponed metavariables. The following
+             example demonstrates the issue fixed by this test.
+             ```
+               structure Point where
+                 x : Nat
+                 y : Nat
+
+               def Point.compute (p : Point) : Point :=
+                 let p := { p with x := 1 }
+                 let p := { p with y := 0 }
+                 if (p.x - p.y) > p.x then p else p
+             ```
+             The `isDefEq` test above fails for `Decidable (p.x - p.y ≤ p.x)` when the structure instance assigned to
+             `p` has not been elaborated yet.
+           -/
+          return false -- we will try again later
+        let oldValType ← inferType oldVal
+        let valType ← inferType val
+        unless (← isDefEq oldValType valType) do
+          let (oldValType, valType) ← addPPExplicitToExposeDiff oldValType valType
+          throwError "synthesized type class instance type is not definitionally equal to expected type, synthesized{indentExpr val}\nhas type{indentExpr valType}\nexpected{indentExpr oldValType}{extraErrorMsg}"
+        let (oldVal, val) ← addPPExplicitToExposeDiff oldVal val
+        throwError "synthesized type class instance is not definitionally equal to expression inferred by typing rules, synthesized{indentExpr val}\ninferred{indentExpr oldVal}{extraErrorMsg}"
+    else
+      unless (← isDefEq (mkMVar instMVar) val) do
+        throwError "failed to assign synthesized type class instance{indentExpr val}{extraErrorMsg}"
+    return true
+  | .undef => return false -- we will try later
+  | .none  =>
+    if (← read).ignoreTCFailures then
+      return false
+    else
+      throwError "failed to synthesize{indentExpr type}{extraErrorMsg}{useDiagnosticMsg}"
+
+def mkCoe (expectedType : Expr) (e : Expr) (f? : Option Expr := none) (errorMsgHeader? : Option String := none)
+    (mkErrorMsg? : Option (MVarId → (expectedType e : Expr) → MetaM MessageData) := none)
+    (mkImmedErrorMsg? : Option ((errorMsg? : Option MessageData) → (expectedType e : Expr) → MetaM MessageData) := none) : TermElabM Expr := do
+  withTraceNode `Elab.coe (fun _ => return m!"adding coercion for {e} : {← inferType e} =?= {expectedType}") do
+  try
+    withoutMacroStackAtErr do
+      match ← coerce? e expectedType with
+      | .some eNew => return eNew
+      | .none => failure
+      | .undef =>
+        let mvarAux ← mkFreshExprMVar expectedType MetavarKind.syntheticOpaque
+        registerSyntheticMVarWithCurrRef mvarAux.mvarId! (.coe errorMsgHeader? expectedType e f? mkErrorMsg?)
+        return mvarAux
+  catch
+    | .error _ msg =>
+      if let some mkImmedErrorMsg := mkImmedErrorMsg? then
+        throwError (← mkImmedErrorMsg msg expectedType e)
+      else
+        throwTypeMismatchError errorMsgHeader? expectedType (← inferType e) e f? msg
+    | _            =>
+      if let some mkImmedErrorMsg := mkImmedErrorMsg? then
+        throwError (← mkImmedErrorMsg none expectedType e)
+      else
+        throwTypeMismatchError errorMsgHeader? expectedType (← inferType e) e f?
+
+def mkCoeWithErrorMsgs (expectedType : Expr) (e : Expr)
+    (mkImmedErrorMsg : (errorMsg? : Option MessageData) → (expectedType e : Expr) → MetaM MessageData)
+    (mkErrorMsg : MVarId → (expectedType e : Expr) → MetaM MessageData) : TermElabM Expr := do
+  mkCoe expectedType e (mkImmedErrorMsg? := mkImmedErrorMsg) (mkErrorMsg? := mkErrorMsg)
+
+/--
+If `expectedType?` is `some t`, then ensures `t` and `eType` are definitionally equal by inserting a coercion if necessary.
+
+Argument `f?` is used only for generating error messages when inserting coercions fails.
+-/
+def ensureHasType (expectedType? : Option Expr) (e : Expr)
+    (errorMsgHeader? : Option String := none) (f? : Option Expr := none) : TermElabM Expr := do
+  let some expectedType := expectedType? | return e
+  if (← isDefEq (← inferType e) expectedType) then
+    return e
+  else
+    mkCoe expectedType e f? errorMsgHeader?
+
+def ensureHasTypeWithErrorMsgs (expectedType? : Option Expr) (e : Expr)
+    (mkImmedErrorMsg : (errorMsg? : Option MessageData) → (expectedType e : Expr) → MetaM MessageData)
+    (mkErrorMsg : MVarId → (expectedType e : Expr) → MetaM MessageData) : TermElabM Expr := do
+  let some expectedType := expectedType? | return e
+  if (← isDefEq (← inferType e) expectedType) then
+    return e
+  else
+    mkCoeWithErrorMsgs expectedType e mkImmedErrorMsg mkErrorMsg
+
+/--
+  Create a synthetic sorry for the given expected type. If `expectedType? = none`, then a fresh
+  metavariable is created to represent the type.
+-/
+private def mkSyntheticSorryFor (expectedType? : Option Expr) : TermElabM Expr := do
+  let expectedType ← match expectedType? with
+    | none              => mkFreshTypeMVar
+    | some expectedType => pure expectedType
+  mkLabeledSorry expectedType (synthetic := true) (unique := false)
+
+/--
+  Log the given exception, and create a synthetic sorry for representing the failed
+  elaboration step with exception `ex`.
+-/
+def exceptionToSorry (ex : Exception) (expectedType? : Option Expr) : TermElabM Expr := do
+  let syntheticSorry ← mkSyntheticSorryFor expectedType?
+  logException ex
+  pure syntheticSorry
+
+/-- If `mayPostpone == true`, throw `Exception.postpone`. -/
+def tryPostpone : TermElabM Unit := do
+  if (← read).mayPostpone then
+    throwPostpone
+
+/-- Return `true` if `e` reduces (by unfolding only `[reducible]` declarations) to `?m ...` -/
+def isMVarApp (e : Expr) : TermElabM Bool :=
+  return (← whnfR e).getAppFn.isMVar
+
+/-- If `mayPostpone == true` and `e`'s head is a metavariable, throw `Exception.postpone`. -/
+def tryPostponeIfMVar (e : Expr) : TermElabM Unit := do
+  if (← isMVarApp e) then
+    tryPostpone
+
+/-- If `e? = some e`, then `tryPostponeIfMVar e`, otherwise it is just `tryPostpone`. -/
+def tryPostponeIfNoneOrMVar (e? : Option Expr) : TermElabM Unit :=
+  match e? with
+  | some e => tryPostponeIfMVar e
+  | none   => tryPostpone
+
+/--
+  Throws `Exception.postpone`, if `expectedType?` contains unassigned metavariables.
+  It is a noop if `mayPostpone == false`.
+-/
+def tryPostponeIfHasMVars? (expectedType? : Option Expr) : TermElabM (Option Expr) := do
+  tryPostponeIfNoneOrMVar expectedType?
+  let some expectedType := expectedType? | return none
+  let expectedType ← instantiateMVars expectedType
+  if expectedType.hasExprMVar then
+    tryPostpone
+    return none
+  return some expectedType
+
+/--
+  Throws `Exception.postpone`, if `expectedType?` contains unassigned metavariables.
+  If `mayPostpone == false`, it throws error `msg`.
+-/
+def tryPostponeIfHasMVars (expectedType? : Option Expr) (msg : String) : TermElabM Expr := do
+  let some expectedType ← tryPostponeIfHasMVars? expectedType? |
+    throwError "{msg}, expected type contains metavariables{indentD expectedType?}"
+  return expectedType
+
+def withExpectedType (expectedType? : Option Expr) (x : Expr → TermElabM Expr) : TermElabM Expr := do
+  tryPostponeIfNoneOrMVar expectedType?
+  let some expectedType ← pure expectedType?
+    | throwError "expected type must be known"
+  x expectedType
+
+/--
+  Save relevant context for term elaboration postponement.
+-/
+def saveContext : TermElabM SavedContext :=
+  return {
+    macroStack := (← read).macroStack
+    declName?  := (← read).declName?
+    options    := (← getOptions)
+    openDecls  := (← getOpenDecls)
+    errToSorry := (← read).errToSorry
+    levelNames := (← get).levelNames
+  }
+
+/--
+  Execute `x` with the context saved using `saveContext`.
+-/
+def withSavedContext (savedCtx : SavedContext) (x : TermElabM α) : TermElabM α := do
+  withReader (fun ctx => { ctx with declName? := savedCtx.declName?, macroStack := savedCtx.macroStack, errToSorry := savedCtx.errToSorry }) <|
+    withTheReader Core.Context (fun ctx => { ctx with options := savedCtx.options, openDecls := savedCtx.openDecls }) <|
+      withLevelNames savedCtx.levelNames x
+
+/--
+Delay the elaboration of `stx`, and return a fresh metavariable that works a placeholder.
+Remark: the caller is responsible for making sure the info tree is properly updated.
+This method is used only at `elabUsingElabFnsAux`.
+-/
+private def postponeElabTermCore (stx : Syntax) (expectedType? : Option Expr) : TermElabM Expr := do
+  trace[Elab.postpone] "{stx} : {expectedType?}"
+  let mvar ← mkFreshExprMVar expectedType? MetavarKind.syntheticOpaque
+  registerSyntheticMVar stx mvar.mvarId! (SyntheticMVarKind.postponed (← saveContext))
+  return mvar
+
+def getSyntheticMVarDecl? (mvarId : MVarId) : TermElabM (Option SyntheticMVarDecl) :=
+  return (← get).syntheticMVars.find? mvarId
+
+register_builtin_option debug.byAsSorry : Bool := {
+  defValue := false
+  group    := "debug"
+  descr    := "replace `by ..` blocks with `sorry` IF the expected type is a proposition"
+}
+
+/--
+Creates a new metavariable of type `type` that will be synthesized using the tactic code.
+The `tacticCode` syntax is the full `by ..` syntax.
+-/
+def mkTacticMVar (type : Expr) (tacticCode : Syntax) (kind : TacticMVarKind)
+    (delayOnMVars := false) : TermElabM Expr := do
+  if ← pure (debug.byAsSorry.get (← getOptions)) <&&> isProp type then
+    withRef tacticCode <| mkLabeledSorry type false (unique := true)
+  else
+    let mvar ← mkFreshExprMVar type MetavarKind.syntheticOpaque
+    let mvarId := mvar.mvarId!
+    let ref ← getRef
+    registerSyntheticMVar ref mvarId <| .tactic tacticCode (← saveContext) kind delayOnMVars
+    return mvar
+
+/--
+  Create an auxiliary annotation to make sure we create an `Info` even if `e` is a metavariable.
+  See `mkTermInfo`.
+
+  We use this function because some elaboration functions elaborate subterms that may not be immediately
+  part of the resulting term. Example:
+  ```
+  let_mvar% ?m := b; wait_if_type_mvar% ?m; body
+  ```
+  If the type of `b` is not known, then `wait_if_type_mvar% ?m; body` is postponed and just returns a fresh
+  metavariable `?n`. The elaborator for
+  ```
+  let_mvar% ?m := b; wait_if_type_mvar% ?m; body
+  ```
+  returns `mkSaveInfoAnnotation ?n` to make sure the info nodes created when elaborating `b` are "saved".
+  This is a bit hackish, but elaborators like `let_mvar%` are rare.
+-/
+def mkSaveInfoAnnotation (e : Expr) : Expr :=
+  if e.isMVar then
+    mkAnnotation `save_info e
+  else
+    e
+
+def isSaveInfoAnnotation? (e : Expr) : Option Expr :=
+  annotation? `save_info e
+
+partial def removeSaveInfoAnnotation (e : Expr) : Expr :=
+  match isSaveInfoAnnotation? e with
+  | some e => removeSaveInfoAnnotation e
+  | _ => e
+
+/--
+  Return `some mvarId` if `e` corresponds to a hole that is going to be filled "later" by executing a tactic or resuming elaboration.
+
+  We do not save `ofTermInfo` for this kind of node in the `InfoTree`.
+-/
+def isTacticOrPostponedHole? (e : Expr) : TermElabM (Option MVarId) := do
+  match e with
+  | Expr.mvar mvarId =>
+    match (← getSyntheticMVarDecl? mvarId) with
+    | some { kind := .tactic .., .. }    => return mvarId
+    | some { kind := .postponed .., .. } => return mvarId
+    | _                                  => return none
+  | _ => pure none
+
+def mkTermInfo (elaborator : Name) (stx : Syntax) (e : Expr) (expectedType? : Option Expr := none)
+    (lctx? : Option LocalContext := none) (isBinder := false) :
+    TermElabM (Sum Info MVarId) := do
+  match (← isTacticOrPostponedHole? e) with
+  | some mvarId => return Sum.inr mvarId
+  | none =>
+    let e := removeSaveInfoAnnotation e
+    return Sum.inl <| Info.ofTermInfo { elaborator, lctx := lctx?.getD (← getLCtx), expr := e, stx, expectedType?, isBinder }
+
+def mkPartialTermInfo (elaborator : Name) (stx : Syntax) (expectedType? : Option Expr := none)
+    (lctx? : Option LocalContext := none) :
+    TermElabM Info := do
+  return Info.ofPartialTermInfo { elaborator, lctx := lctx?.getD (← getLCtx), stx, expectedType? }
+
+/--
+Pushes a new leaf node to the info tree associating the expression `e` to the syntax `stx`.
+As a result, when the user hovers over `stx` they will see the type of `e`, and if `e`
+is a constant they will see the constant's doc string.
+
+* `expectedType?`: the expected type of `e` at the point of elaboration, if available
+* `lctx?`: the local context in which to interpret `e` (otherwise it will use `← getLCtx`)
+* `elaborator`: a declaration name used as an alternative target for go-to-definition
+* `isBinder`: if true, this will be treated as defining `e` (which should be a local constant)
+  for the purpose of go-to-definition on local variables
+* `force`: In patterns, the effect of `addTermInfo` is usually suppressed and replaced
+  by a `patternWithRef?` annotation which will be turned into a term info on the
+  post-match-elaboration expression. This flag overrides that behavior and adds the term
+  info immediately. (See https://github.com/leanprover/lean4/pull/1664.)
+-/
+def addTermInfo (stx : Syntax) (e : Expr) (expectedType? : Option Expr := none)
+    (lctx? : Option LocalContext := none) (elaborator := Name.anonymous)
+    (isBinder := false) (force := false) : TermElabM Expr := do
+  if (← read).inPattern && !force then
+    return mkPatternWithRef e stx
+  else
+    discard <| withInfoContext'
+      (pure ())
+      (fun _ => mkTermInfo elaborator stx e expectedType? lctx? isBinder)
+      (mkPartialTermInfo elaborator stx expectedType? lctx?)
+    return e
+
+def addTermInfo' (stx : Syntax) (e : Expr) (expectedType? : Option Expr := none) (lctx? : Option LocalContext := none) (elaborator := Name.anonymous) (isBinder := false) : TermElabM Unit :=
+  discard <| addTermInfo stx e expectedType? lctx? elaborator isBinder
+
+def withInfoContext' (stx : Syntax) (x : TermElabM Expr)
+    (mkInfo : Expr → TermElabM (Sum Info MVarId)) (mkInfoOnError : TermElabM Info) :
+    TermElabM Expr := do
+  if (← read).inPattern then
+    let e ← x
+    return mkPatternWithRef e stx
+  else
+    Elab.withInfoContext' x mkInfo mkInfoOnError
+
+/-- Info node capturing `def/let rec` bodies, used by the unused variables linter. -/
+structure BodyInfo where
+  /-- The body as a fully elaborated term. `none` if the body failed to elaborate. -/
+  value? : Option Expr
+deriving TypeName
+
+/-- Creates an `Info.ofCustomInfo` node backed by a `BodyInfo`. -/
+def mkBodyInfo (stx : Syntax) (value? : Option Expr) : Info :=
+  .ofCustomInfo { stx, value := .mk { value? : BodyInfo } }
+
+/-- Extracts a `BodyInfo` custom info. -/
+def getBodyInfo? : Info → Option BodyInfo
+  | .ofCustomInfo { value, .. } => value.get? BodyInfo
+  | _ => none
+
+def withTermInfoContext' (elaborator : Name) (stx : Syntax) (x : TermElabM Expr)
+    (expectedType? : Option Expr := none) (lctx? : Option LocalContext := none)
+    (isBinder : Bool := false) :
+    TermElabM Expr :=
+  withInfoContext' stx x
+    (mkTermInfo elaborator stx (expectedType? := expectedType?) (lctx? := lctx?) (isBinder := isBinder))
+    (mkPartialTermInfo elaborator stx (expectedType? := expectedType?) (lctx? := lctx?))
+
+/--
+Postpone the elaboration of `stx`, return a metavariable that acts as a placeholder, and
+ensures the info tree is updated and a hole id is introduced.
+When `stx` is elaborated, new info nodes are created and attached to the new hole id in the info tree.
+-/
+def postponeElabTerm (stx : Syntax) (expectedType? : Option Expr) : TermElabM Expr := do
+  withTermInfoContext' .anonymous stx (expectedType? := expectedType?) do
+    postponeElabTermCore stx expectedType?
+
+/--
+  Helper function for `elabTerm` that tries the registered elaboration functions for `stxNode` kind until it finds one that supports the syntax or
+  an error is found. -/
+private def elabUsingElabFnsAux (s : SavedState) (stx : Syntax) (expectedType? : Option Expr) (catchExPostpone : Bool)
+    : List (KeyedDeclsAttribute.AttributeEntry TermElab) → TermElabM Expr
+  | []                => do throwError "unexpected syntax{indentD stx}"
+  | (elabFn::elabFns) =>
+    try
+      -- record elaborator in info tree, but only when not backtracking to other elaborators (outer `try`)
+      withTermInfoContext' elabFn.declName stx (expectedType? := expectedType?)
+        (try
+          elabFn.value stx expectedType?
+        catch ex => match ex with
+          | .error .. =>
+            if (← read).errToSorry then
+              exceptionToSorry ex expectedType?
+            else
+              throw ex
+          | .internal id _ =>
+            if (← read).errToSorry && id == abortTermExceptionId then
+              exceptionToSorry ex expectedType?
+            else if id == unsupportedSyntaxExceptionId then
+              throw ex  -- to outer try
+            else if catchExPostpone && id == postponeExceptionId then
+              /- If `elab` threw `Exception.postpone`, we reset any state modifications.
+                For example, we want to make sure pending synthetic metavariables created by `elab` before
+                it threw `Exception.postpone` are discarded.
+                Note that we are also discarding the messages created by `elab`.
+
+                For example, consider the expression.
+                `((f.x a1).x a2).x a3`
+                Now, suppose the elaboration of `f.x a1` produces an `Exception.postpone`.
+                Then, a new metavariable `?m` is created. Then, `?m.x a2` also throws `Exception.postpone`
+                because the type of `?m` is not yet known. Then another, metavariable `?n` is created, and
+                finally `?n.x a3` also throws `Exception.postpone`. If we did not restore the state, we would
+                keep "dead" metavariables `?m` and `?n` on the pending synthetic metavariable list. This is
+                wasteful because when we resume the elaboration of `((f.x a1).x a2).x a3`, we start it from scratch
+                and new metavariables are created for the nested functions. -/
+              s.restore
+              postponeElabTermCore stx expectedType?
+            else
+              throw ex)
+    catch ex => match ex with
+      | .internal id _ =>
+        if id == unsupportedSyntaxExceptionId then
+          s.restore  -- also removes the info tree created above
+          elabUsingElabFnsAux s stx expectedType? catchExPostpone elabFns
+        else
+          throw ex
+      | _ => throw ex
+
+private def elabUsingElabFns (stx : Syntax) (expectedType? : Option Expr) (catchExPostpone : Bool) : TermElabM Expr := do
+  let s ← saveState
+  let k := stx.getKind
+  match termElabAttribute.getEntries (← getEnv) k with
+  | []      => throwError "elaboration function for '{k}' has not been implemented{indentD stx}"
+  | elabFns => elabUsingElabFnsAux s stx expectedType? catchExPostpone elabFns
+
+instance : MonadMacroAdapter TermElabM where
+  getCurrMacroScope := getCurrMacroScope
+  getNextMacroScope := return (← getThe Core.State).nextMacroScope
+  setNextMacroScope next := modifyThe Core.State fun s => { s with nextMacroScope := next }
+
+private def isExplicit (stx : Syntax) : Bool :=
+  match stx with
+  | `(@$_) => true
+  | _      => false
+
+private def isExplicitApp (stx : Syntax) : Bool :=
+  stx.getKind == ``Lean.Parser.Term.app && isExplicit stx[0]
+
+/--
+  Return true if `stx` is a lambda abstraction containing a `{}` or `[]` binder annotation.
+  Example: `fun {α} (a : α) => a` -/
+private def isLambdaWithImplicit (stx : Syntax) : Bool :=
+  match stx with
+  | `(fun $binders* => $_) => binders.raw.any fun b => b.isOfKind ``Lean.Parser.Term.implicitBinder || b.isOfKind `Lean.Parser.Term.instBinder
+  | _                      => false
+
+private partial def dropTermParens : Syntax → Syntax := fun stx =>
+  match stx with
+  | `(($stx)) => dropTermParens stx
+  | _         => stx
+
+private def isHole (stx : Syntax) : Bool :=
+  stx.isOfKind ``Lean.Parser.Term.hole || stx.isOfKind ``Lean.Parser.Term.syntheticHole
+
+private def isTacticBlock (stx : Syntax) : Bool :=
+  match stx with
+  | `(by $_:tacticSeq) => true
+  | _ => false
+
+private def isNoImplicitLambda (stx : Syntax) : Bool :=
+  match stx with
+  | `(no_implicit_lambda% $_:term) => true
+  | _ => false
+
+private def isTypeAscription (stx : Syntax) : Bool :=
+  match stx with
+  | `(($_ : $[$_]?)) => true
+  | _                => false
+
+def hasNoImplicitLambdaAnnotation (type : Expr) : Bool :=
+  annotation? `noImplicitLambda type |>.isSome
+
+def mkNoImplicitLambdaAnnotation (type : Expr) : Expr :=
+  if hasNoImplicitLambdaAnnotation type then
+    type
+  else
+    mkAnnotation `noImplicitLambda type
+
+/-- Block usage of implicit lambdas if `stx` is `@f` or `@f arg1 ...` or `fun` with an implicit binder annotation. -/
+def blockImplicitLambda (stx : Syntax) : Bool :=
+  let stx := dropTermParens stx
+  -- TODO: make it extensible
+  isExplicit stx || isExplicitApp stx || isLambdaWithImplicit stx || isHole stx || isTacticBlock stx ||
+  isNoImplicitLambda stx || isTypeAscription stx
+
+/-- Return true iff `stx` is a `Syntax.ident`, and it is a local variable. -/
+def isLocalIdent? (stx : Syntax) : TermElabM (Option Expr) :=
+  match stx with
+  | Syntax.ident _ _ val _ => do
+    let r? ← resolveLocalName val
+    match r? with
+    | some (fvar, []) => return some fvar
+    | _               => return none
+  | _ => return none
+
+inductive UseImplicitLambdaResult where
+  | no
+  | yes (expectedType : Expr)
+  | postpone
+
+/--
+  Return normalized expected type if it is of the form `{a : α} → β` or `[a : α] → β` and
+  `blockImplicitLambda stx` is not true, else return `none`.
+
+  Remark: implicit lambdas are not triggered by the strict implicit binder annotation `{{a : α}} → β`
+-/
+private def useImplicitLambda (stx : Syntax) (expectedType? : Option Expr) : TermElabM UseImplicitLambdaResult := do
+  if blockImplicitLambda stx then
+    return .no
+  let some expectedType := expectedType? | return .no
+  if hasNoImplicitLambdaAnnotation expectedType then
+    return .no
+  let expectedType ← whnfForall expectedType
+  let .forallE _ _ _ c := expectedType | return .no
+  unless c.isImplicit || c.isInstImplicit do
+    return .no
+  if let some x ← isLocalIdent? stx then
+    if (← isMVarApp (← inferType x)) then
+      /-
+      If `stx` is a local variable without type information, then adding implicit lambdas makes elaboration fail.
+      We should try to postpone elaboration until the type of the local variable becomes available, or disable
+      implicit lambdas if we cannot postpone anymore.
+      Here is an example where this special case is useful.
+      ```
+      def foo2mk (_ : ∀ {α : Type} (a : α), a = a) : nat := 37
+      example (x) : foo2mk x = foo2mk x := rfl
+      ```
+      The example about would fail without this special case.
+      The expected type would be `(a : α✝) → a = a`, where `α✝` is a new free variable introduced by the implicit lambda.
+      Now, let `?m` be the type of `x`. Then, the constraint `?m =?= (a : α✝) → a = a` cannot be solved using the
+      assignment `?m := (a : α✝) → a = a` since `α✝` is not in the scope of `?m`.
+
+      Note that, this workaround does not prevent the following example from failing.
+      ```
+      example (x) : foo2mk (id x) = 37 := rfl
+      ```
+      The user can write
+      ```
+      example (x) : foo2mk (id @x) = 37 := rfl
+      ```
+      -/
+      return .postpone
+  return .yes expectedType
+
+private def decorateErrorMessageWithLambdaImplicitVars (ex : Exception) (impFVars : Array Expr) : TermElabM Exception := do
+  match ex with
+  | .error ref msg =>
+    if impFVars.isEmpty then
+      return Exception.error ref msg
+    else
+      let mut msg := m!"{msg}\nthe following variables have been introduced by the implicit lambda feature"
+      for impFVar in impFVars do
+        let auxMsg := m!"{impFVar} : {← inferType impFVar}"
+        let auxMsg ← addMessageContext auxMsg
+        msg := m!"{msg}{indentD auxMsg}"
+      msg := m!"{msg}\nyou can disable implicit lambdas using `@` or writing a lambda expression with `\{}` or `[]` binder annotations."
+      return Exception.error ref msg
+  | _ => return ex
+
+private def elabImplicitLambdaAux (stx : Syntax) (catchExPostpone : Bool) (expectedType : Expr) (impFVars : Array Expr) : TermElabM Expr := do
+  let body ← elabUsingElabFns stx expectedType catchExPostpone
+  try
+    let body ← ensureHasType expectedType body
+    let r ← mkLambdaFVars impFVars body
+    trace[Elab.implicitForall] r
+    return r
+  catch ex =>
+    throw (← decorateErrorMessageWithLambdaImplicitVars ex impFVars)
+
+private partial def elabImplicitLambda (stx : Syntax) (catchExPostpone : Bool) (type : Expr) : TermElabM Expr :=
+  loop type #[]
+where
+  loop (type : Expr) (fvars : Array Expr) : TermElabM Expr := do
+    match (← whnfForall type) with
+    | .forallE n d b c =>
+      if c.isExplicit then
+        elabImplicitLambdaAux stx catchExPostpone type fvars
+      else withFreshMacroScope do
+        let n ← MonadQuotation.addMacroScope n
+        withLocalDecl n c d fun fvar => do
+          let type := b.instantiate1 fvar
+          loop type (fvars.push fvar)
+    | _ =>
+      elabImplicitLambdaAux stx catchExPostpone type fvars
+
+/-- Main loop for `elabTerm` -/
+private partial def elabTermAux (expectedType? : Option Expr) (catchExPostpone : Bool) (implicitLambda : Bool) : Syntax → TermElabM Expr
+  | .missing => mkSyntheticSorryFor expectedType?
+  | stx => withFreshMacroScope <| withIncRecDepth do
+    withTraceNode `Elab.step (fun _ => return m!"expected type: {expectedType?}, term\n{stx}")
+      (tag := stx.getKind.toString) do
+    checkSystem "elaborator"
+    let env ← getEnv
+    let result ← match (← liftMacroM (expandMacroImpl? env stx)) with
+    | some (decl, stxNew?) =>
+      let stxNew ← liftMacroM <| liftExcept stxNew?
+      withTermInfoContext' decl stx (expectedType? := expectedType?) <|
+        withMacroExpansion stx stxNew <|
+          withRef stxNew <|
+            elabTermAux expectedType? catchExPostpone implicitLambda stxNew
+    | _ =>
+      let useImplicitResult ← if implicitLambda && (← read).implicitLambda then useImplicitLambda stx expectedType? else pure .no
+      match useImplicitResult with
+      | .yes expectedType => elabImplicitLambda stx catchExPostpone expectedType
+      | .no => elabUsingElabFns stx expectedType? catchExPostpone
+      | .postpone =>
+        /-
+        Try to postpone elaboration, and if we cannot postpone anymore disable implicit lambdas.
+        See comment at `useImplicitLambda`.
+        -/
+        if (← read).mayPostpone then
+          if catchExPostpone then
+            postponeElabTerm stx expectedType?
+          else
+            throwPostpone
+        else
+          elabUsingElabFns stx expectedType? catchExPostpone
+    trace[Elab.step.result] result
+    pure result
+
+/-- Store in the `InfoTree` that `e` is a "dot"-completion target. `stx` should cover the entire term. -/
+def addDotCompletionInfo (stx : Syntax) (e : Expr) (expectedType? : Option Expr) : TermElabM Unit := do
+  addCompletionInfo <| CompletionInfo.dot { expr := e, stx, lctx := (← getLCtx), elaborator := .anonymous, expectedType? } (expectedType? := expectedType?)
+
+/--
+  Main function for elaborating terms.
+  It extracts the elaboration methods from the environment using the node kind.
+  Recall that the environment has a mapping from `SyntaxNodeKind` to `TermElab` methods.
+  It creates a fresh macro scope for executing the elaboration method.
+  All unlogged trace messages produced by the elaboration method are logged using
+  the position information at `stx`. If the elaboration method throws an `Exception.error` and `errToSorry == true`,
+  the error is logged and a synthetic sorry expression is returned.
+  If the elaboration throws `Exception.postpone` and `catchExPostpone == true`,
+  a new synthetic metavariable of kind `SyntheticMVarKind.postponed` is created, registered,
+  and returned.
+  The option `catchExPostpone == false` is used to implement `resumeElabTerm`
+  to prevent the creation of another synthetic metavariable when resuming the elaboration.
+
+  If `implicitLambda == false`, then disable implicit lambdas feature for the given syntax, but not for its subterms.
+  We use this flag to implement, for example, the `@` modifier. If `Context.implicitLambda == false`, then this parameter has no effect.
+  -/
+def elabTerm (stx : Syntax) (expectedType? : Option Expr) (catchExPostpone := true) (implicitLambda := true) : TermElabM Expr :=
+  withRef stx <| elabTermAux expectedType? catchExPostpone implicitLambda stx
+
+/--
+Similar to `Lean.Elab.Term.elabTerm`, but ensures that the type of the elaborated term is `expectedType?`
+by inserting coercions if necessary.
+
+If `errToSorry` is true, then if coercion insertion fails, this function returns `sorry` and logs the error.
+Otherwise, it throws the error.
+-/
+def elabTermEnsuringType (stx : Syntax) (expectedType? : Option Expr) (catchExPostpone := true) (implicitLambda := true) (errorMsgHeader? : Option String := none) : TermElabM Expr := do
+  let e ← elabTerm stx expectedType? catchExPostpone implicitLambda
+  try
+    withRef stx <| ensureHasType expectedType? e errorMsgHeader?
+  catch ex =>
+    if (← read).errToSorry && ex matches .error .. then
+      withRef stx <| exceptionToSorry ex expectedType?
+    else
+      throw ex
+
+/-- Execute `x` and return `some` if no new errors were recorded or exceptions were thrown. Otherwise, return `none`. -/
+def commitIfNoErrors? (x : TermElabM α) : TermElabM (Option α) := do
+  let saved ← saveState
+  Core.resetMessageLog
+  try
+    let a ← x
+    if (← MonadLog.hasErrors) then
+      restoreState saved
+      return none
+    else
+      Core.setMessageLog (saved.meta.core.messages ++ (← Core.getMessageLog))
+      return a
+  catch _ =>
+    restoreState saved
+    return none
+
+/-- Adapt a syntax transformation to a regular, term-producing elaborator. -/
+def adaptExpander (exp : Syntax → TermElabM Syntax) : TermElab := fun stx expectedType? => do
+  let stx' ← exp stx
+  withMacroExpansion stx stx' <| elabTerm stx' expectedType?
+
+/--
+  Create a new metavariable with the given type, and try to synthesize it.
+  If type class resolution cannot be executed (e.g., it is stuck because of metavariables in `type`),
+  register metavariable as a pending one.
+-/
+def mkInstMVar (type : Expr) (extraErrorMsg? : Option MessageData := none) : TermElabM Expr := do
+  let mvar ← mkFreshExprMVar type MetavarKind.synthetic
+  let mvarId := mvar.mvarId!
+  unless (← synthesizeInstMVarCore mvarId (extraErrorMsg? := extraErrorMsg?)) do
+    registerSyntheticMVarWithCurrRef mvarId (.typeClass extraErrorMsg?)
+  return mvar
+
+/--
+  Make sure `e` is a type by inferring its type and making sure it is an `Expr.sort`
+  or is unifiable with `Expr.sort`, or can be coerced into one. -/
+def ensureType (e : Expr) : TermElabM Expr := do
+  if (← isType e) then
+    return e
+  else
+    let eType ← inferType e
+    let u ← mkFreshLevelMVar
+    if (← isDefEq eType (mkSort u)) then
+      return e
+    else if let some coerced ← coerceToSort? e then
+      return coerced
+    else
+      if (← instantiateMVars e).hasSyntheticSorry then
+        throwAbortTerm
+      throwError "type expected, got\n  ({← instantiateMVars e} : {← instantiateMVars eType})"
+
+/-- Elaborate `stx` and ensure result is a type. -/
+def elabType (stx : Syntax) : TermElabM Expr := do
+  let u ← mkFreshLevelMVar
+  let type ← elabTerm stx (mkSort u)
+  withRef stx <| ensureType type
+
+/--
+  Enable auto-bound implicits, and execute `k` while catching auto bound implicit exceptions. When an exception is caught,
+  a new local declaration is created, registered, and `k` is tried to be executed again. -/
+partial def withAutoBoundImplicit (k : TermElabM α) : TermElabM α := do
+  let flag := autoImplicit.get (← getOptions)
+  if flag then
+    withReader (fun ctx => { ctx with autoBoundImplicit := flag, autoBoundImplicits := {} }) do
+      let rec loop (s : SavedState) : TermElabM α := withIncRecDepth do
+        checkSystem "auto-implicit"
+        try
+          k
+        catch
+          | ex => match isAutoBoundImplicitLocalException? ex with
+            | some n =>
+              -- Restore state, declare `n`, and try again
+              s.restore (restoreInfo := true)
+              withLocalDecl n .implicit (← mkFreshTypeMVar) fun x =>
+                withReader (fun ctx => { ctx with autoBoundImplicits := ctx.autoBoundImplicits.push x } ) do
+                  loop (← saveState)
+            | none   => throw ex
+      loop (← saveState)
+  else
+    k
+
+def withoutAutoBoundImplicit (k : TermElabM α) : TermElabM α := do
+  withReader (fun ctx => { ctx with autoBoundImplicit := false, autoBoundImplicits := {} }) k
+
+partial def withAutoBoundImplicitForbiddenPred (p : Name → Bool) (x : TermElabM α) : TermElabM α := do
+  withReader (fun ctx => { ctx with autoBoundImplicitForbidden := fun n => p n || ctx.autoBoundImplicitForbidden n }) x
+
+/--
+  Collect unassigned metavariables in `type` that are not already in `init` and not satisfying `except`.
+-/
+partial def collectUnassignedMVars (type : Expr) (init : Array Expr := #[]) (except : MVarId → Bool := fun _ => false)
+    : TermElabM (Array Expr) := do
+  let mvarIds ← getMVars type
+  if mvarIds.isEmpty then
+    return init
+  else
+    go mvarIds.toList init init
+where
+  go (mvarIds : List MVarId) (result visited : Array Expr) : TermElabM (Array Expr) := do
+    match mvarIds with
+    | [] => return result
+    | mvarId :: mvarIds => do
+      let visited := visited.push (mkMVar mvarId)
+      if (← mvarId.isAssigned) then
+        go mvarIds result visited
+      else if result.contains (mkMVar mvarId) || except mvarId then
+        go mvarIds result visited
+      else
+        let mvarType := (← getMVarDecl mvarId).type
+        let mvarIdsNew ← getMVars mvarType
+        let mvarIdsNew := mvarIdsNew.filter fun mvarId => !visited.contains (mkMVar mvarId)
+        if mvarIdsNew.isEmpty then
+          go mvarIds (result.push (mkMVar mvarId)) visited
+        else
+          go (mvarIdsNew.toList ++ mvarId :: mvarIds) result visited
+
+/--
+Adds an `InlayHintInfo` for the fvar auto implicits in `autos` at `inlayHintPos`.
+The inserted inlay hint has a hover that denotes the type of the auto-implicit (with meta-variables)
+and can be inserted at `inlayHintPos`.
+-/
+def addAutoBoundImplicitsInlayHint (autos : Array Expr) (inlayHintPos : String.Pos) : TermElabM Unit := do
+  -- If the list of auto-implicits contains a non-type fvar, then the list of auto-implicits will
+  -- also contain an mvar that denotes the type of the non-type fvar.
+  -- For example, the auto-implicit `x` in a type `Foo x` for `Foo.{u} {α : Sort u} (x : α) : Type`
+  -- also comes with an auto-implicit mvar denoting the type of `x`.
+  -- We have no way of displaying this mvar to the user in an inlay hint, as it doesn't have a name,
+  -- so we filter it.
+  -- This also means that inserting the inlay hint with the syntax displayed in the inlay hint will
+  -- cause a "failed to infer binder type" error, since we don't have a name to insert in the code.
+  let autos := autos.filter (· matches .fvar ..)
+  if autos.isEmpty then
+    return
+  let autoNames ← autos.mapM (·.fvarId!.getUserName)
+  let formattedHint := s!" \{{" ".intercalate <| Array.toList <| autoNames.map toString}}"
+  let deferredResolution ih := do
+    let description := "Automatically-inserted implicit parameters:"
+    let codeBlockStart := "```lean"
+    let typeInfos ← autos.mapM fun auto => do
+      let name := toString <| ← auto.fvarId!.getUserName
+      let type := toString <| ← Meta.ppExpr <| ← instantiateMVars (← inferType auto)
+      return s!"{name} : {type}"
+    let codeBlockEnd := "```"
+    let tooltip := "\n".intercalate <| description :: codeBlockStart :: typeInfos.toList ++ [codeBlockEnd]
+    return { ih with tooltip? := tooltip }
+  pushInfoLeaf <| .ofCustomInfo {
+      position := inlayHintPos
+      label := .name formattedHint
+      textEdits := #[{
+        range := ⟨inlayHintPos, inlayHintPos⟩,
+        newText := formattedHint
+      }]
+      kind? := some .parameter
+      lctx := ← getLCtx
+      deferredResolution
+      : InlayHint
+    }.toCustomInfo
+
+/--
+  Return `autoBoundImplicits ++ xs`
+  This method throws an error if a variable in `autoBoundImplicits` depends on some `x` in `xs`.
+  The `autoBoundImplicits` may contain free variables created by the auto-implicit feature, and unassigned free variables.
+  It avoids the hack used at `autoBoundImplicitsOld`.
+
+  If `inlayHintPos?` is set, this function also inserts an inlay hint denoting `autoBoundImplicits`.
+  See `addAutoBoundImplicitsInlayHint` for more information.
+
+  Remark: we cannot simply replace every occurrence of `addAutoBoundImplicitsOld` with this one because a particular
+  use-case may not be able to handle the metavariables in the array being given to `k`.
+-/
+def addAutoBoundImplicits (xs : Array Expr) (inlayHintPos? : Option String.Pos) : TermElabM (Array Expr) := do
+  let autos := (← read).autoBoundImplicits
+  go autos.toList #[]
+where
+  go (todo : List Expr) (autos : Array Expr) : TermElabM (Array Expr) := do
+    match todo with
+    | [] =>
+      if let some inlayHintPos := inlayHintPos? then
+        addAutoBoundImplicitsInlayHint autos inlayHintPos
+      for auto in autos do
+        if auto.isFVar then
+          let localDecl ← auto.fvarId!.getDecl
+          for x in xs do
+            if (← localDeclDependsOn localDecl x.fvarId!) then
+              throwError "invalid auto implicit argument '{auto}', it depends on explicitly provided argument '{x}'"
+      return autos ++ xs
+    | auto :: todo =>
+      let autos ← collectUnassignedMVars (← inferType auto) autos
+      go todo (autos.push auto)
+
+/--
+  Similar to `addAutoBoundImplicits`, but converts all metavariables into free variables.
+
+  It uses `mkForallFVars` + `forallBoundedTelescope` to convert metavariables into free variables.
+  The type `type` is modified during the process if type depends on `xs`.
+  We use this method to simplify the conversion of code using `autoBoundImplicitsOld` to `autoBoundImplicits`.
+-/
+def addAutoBoundImplicits' (xs : Array Expr) (type : Expr) (k : Array Expr → Expr → TermElabM α) (inlayHintPos? : Option String.Pos := none) : TermElabM α := do
+  let xs ← addAutoBoundImplicits xs inlayHintPos?
+  if xs.all (·.isFVar) then
+    k xs type
+  else
+    forallBoundedTelescope (← mkForallFVars xs type) xs.size fun xs type => k xs type
+
+def mkAuxName (suffix : Name) : TermElabM Name := mkAuxDeclName (kind := suffix)
+
+builtin_initialize registerTraceClass `Elab.letrec
+
+/-- Return true if mvarId is an auxiliary metavariable created for compiling `let rec` or it
+   is delayed assigned to one. -/
+def isLetRecAuxMVar (mvarId : MVarId) : TermElabM Bool := do
+  trace[Elab.letrec] "mvarId: {mkMVar mvarId} letrecMVars: {(← get).letRecsToLift.map (mkMVar $ ·.mvarId)}"
+  let mvarId ← getDelayedMVarRoot mvarId
+  trace[Elab.letrec] "mvarId root: {mkMVar mvarId}"
+  return (← get).letRecsToLift.any (·.mvarId == mvarId)
+
+private def checkDeprecatedCore (constName : Name) : TermElabM Unit := do
+  if (← read).checkDeprecated then
+    Linter.checkDeprecated constName
+
+/--
+  Create an `Expr.const` using the given name and explicit levels.
+  Remark: fresh universe metavariables are created if the constant has more universe
+  parameters than `explicitLevels`.
+
+  If `checkDeprecated := true`, then `Linter.checkDeprecated` is invoked.
+-/
+def mkConst (constName : Name) (explicitLevels : List Level := []) : TermElabM Expr := do
+  checkDeprecatedCore constName
+  let cinfo ← getConstVal constName
+  if explicitLevels.length > cinfo.levelParams.length then
+    throwError "too many explicit universe levels for '{constName}'"
+  else
+    let numMissingLevels := cinfo.levelParams.length - explicitLevels.length
+    let us ← mkFreshLevelMVars numMissingLevels
+    return Lean.mkConst constName (explicitLevels ++ us)
+
+def checkDeprecated (ref : Syntax) (e : Expr) : TermElabM Unit := do
+  if let .const declName _ := e.getAppFn then
+    withRef ref do checkDeprecatedCore declName
+
+@[inline] def withoutCheckDeprecated [MonadWithReaderOf Context m] : m α → m α :=
+  withTheReader Context (fun ctx => { ctx with checkDeprecated := false })
+
+private def mkConsts (candidates : List (Name × List String)) (explicitLevels : List Level) : TermElabM (List (Expr × List String)) := do
+  candidates.foldlM (init := []) fun result (declName, projs) => do
+    -- TODO: better support for `mkConst` failure. We may want to cache the failures, and report them if all candidates fail.
+    /-
+    We disable `checkDeprecated` here because there may be many overloaded symbols.
+    Note that, this method and `resolveName` and `resolveName'` return a list of pairs instead of a list of `TermElabResult`s.
+    We perform the `checkDeprecated` test at `resolveId?` and `elabAppFnId`.
+    At `elabAppFnId`, we perform the check when converting the list returned by `resolveName'` into a list of
+    `TermElabResult`s.
+    -/
+    let const ← withoutCheckDeprecated <| mkConst declName explicitLevels
+    return (const, projs) :: result
+
+def resolveName (stx : Syntax) (n : Name) (preresolved : List Syntax.Preresolved) (explicitLevels : List Level) (expectedType? : Option Expr := none) : TermElabM (List (Expr × List String)) := do
+  addCompletionInfo <| CompletionInfo.id stx stx.getId (danglingDot := false) (← getLCtx) expectedType?
+  if let some (e, projs) ← resolveLocalName n then
+    unless explicitLevels.isEmpty do
+      throwError "invalid use of explicit universe parameters, '{e}' is a local variable"
+    return [(e, projs)]
+  let preresolved := preresolved.filterMap fun
+    | .decl n projs => some (n, projs)
+    | _             => none
+  -- check for section variable capture by a quotation
+  let ctx ← read
+  if let some (e, projs) := preresolved.findSome? fun (n, projs) => ctx.sectionFVars.find? n |>.map (·, projs) then
+    return [(e, projs)]  -- section variables should shadow global decls
+  if preresolved.isEmpty then
+    process (← realizeGlobalName n)
+  else
+    process preresolved
+where
+  process (candidates : List (Name × List String)) : TermElabM (List (Expr × List String)) := do
+    if candidates.isEmpty then
+      if (← read).autoBoundImplicit &&
+           !(← read).autoBoundImplicitForbidden n &&
+           isValidAutoBoundImplicitName n (relaxedAutoImplicit.get (← getOptions)) then
+        throwAutoBoundImplicitLocal n
+      else
+        throwUnknownIdentifierAt stx m!"unknown identifier '{Lean.mkConst n}'"
+    mkConsts candidates explicitLevels
+
+/--
+  Similar to `resolveName`, but creates identifiers for the main part and each projection with position information derived from `ident`.
+  Example: Assume resolveName `v.head.bla.boo` produces `(v.head, ["bla", "boo"])`, then this method produces
+  `(v.head, id, [f₁, f₂])` where `id` is an identifier for `v.head`, and `f₁` and `f₂` are identifiers for fields `"bla"` and `"boo"`. -/
+def resolveName' (ident : Syntax) (explicitLevels : List Level) (expectedType? : Option Expr := none) : TermElabM (List (Expr × Syntax × List Syntax)) := do
+  match ident with
+  | .ident _ _ n preresolved =>
+    let r ← resolveName ident n preresolved explicitLevels expectedType?
+    r.mapM fun (c, fields) => do
+      let ids := ident.identComponents (nFields? := fields.length)
+      return (c, ids.head!, ids.tail!)
+  | _ => throwError "identifier expected"
+
+def resolveId? (stx : Syntax) (kind := "term") (withInfo := false) : TermElabM (Option Expr) := withRef stx do
+  match stx with
+  | .ident _ _ val preresolved =>
+    let rs ← try resolveName stx val preresolved [] catch _ => pure []
+    let rs := rs.filter fun ⟨_, projs⟩ => projs.isEmpty
+    let fs := rs.map fun (f, _) => f
+    match fs with
+    | []  => return none
+    | [f] =>
+      let f ← if withInfo then addTermInfo stx f else pure f
+      checkDeprecated stx f
+      return some f
+    | _   => throwError "ambiguous {kind}, use fully qualified name, possible interpretations {fs}"
+  | _ => throwError "identifier expected"
+
+def TermElabM.run (x : TermElabM α) (ctx : Context := {}) (s : State := {}) : MetaM (α × State) :=
+  withConfig setElabConfig (x ctx |>.run s)
+
+@[inline] def TermElabM.run' (x : TermElabM α) (ctx : Context := {}) (s : State := {}) : MetaM α :=
+  (·.1) <$> x.run ctx s
+
+def TermElabM.toIO (x : TermElabM α)
+    (ctxCore : Core.Context) (sCore : Core.State)
+    (ctxMeta : Meta.Context) (sMeta : Meta.State)
+    (ctx : Context) (s : State) : IO (α × Core.State × Meta.State × State) := do
+  let ((a, s), sCore, sMeta) ← (x.run ctx s).toIO ctxCore sCore ctxMeta sMeta
+  return (a, sCore, sMeta, s)
+
+/--
+  Execute `x` and then tries to solve pending universe constraints.
+  Note that, stuck constraints will not be discarded.
+-/
+def universeConstraintsCheckpoint (x : TermElabM α) : TermElabM α := do
+  let a ← x
+  discard <| processPostponed (mayPostpone := true) (exceptionOnFailure := true)
+  return a
+
+def expandDeclId (currNamespace : Name) (currLevelNames : List Name) (declId : Syntax) (modifiers : Modifiers) : TermElabM ExpandDeclIdResult := do
+  let r ← Elab.expandDeclId currNamespace currLevelNames declId modifiers
+  if (← read).sectionVars.contains r.shortName then
+    throwError "invalid declaration name '{r.shortName}', there is a section variable with the same name"
+  return r
+
+/--
+  Helper function for "embedding" an `Expr` in `Syntax`.
+  It creates a named hole `?m` and immediately assigns `e` to it.
+  Examples:
+  ```lean
+  let e := mkConst ``Nat.zero
+  `(Nat.succ $(← exprToSyntax e))
+  ```
+-/
+def exprToSyntax (e : Expr) : TermElabM Term := withFreshMacroScope do
+  let result ← `(?m)
+  let eType ← inferType e
+  let mvar ← elabTerm result eType
+  mvar.mvarId!.assign e
+  return result
+
+end Term
+
+open Term in
+def withoutModifyingStateWithInfoAndMessages [MonadControlT TermElabM m] [Monad m] (x : m α) : m α := do
+  controlAt TermElabM fun runInBase => withoutModifyingStateWithInfoAndMessagesImpl <| runInBase x
+
+builtin_initialize
+  registerTraceClass `Elab.postpone
+  registerTraceClass `Elab.coe
+  registerTraceClass `Elab.debug
+  registerTraceClass `Elab.reuse
+
+/--
+Marks an elaborator (tactic or command, currently) as supporting incremental elaboration.
+
+For unmarked elaborators, the corresponding snapshot bundle field in the elaboration context is
+unset so as to prevent accidental, incorrect reuse.
+-/
+@[builtin_doc]
+builtin_initialize incrementalAttr : TagAttribute ←
+  registerTagAttribute `incremental "Marks an elaborator (tactic or command, currently) as \
+supporting incremental elaboration. For unmarked elaborators, the corresponding snapshot bundle \
+field in the elaboration context is unset so as to prevent accidental, incorrect reuse."
+
+builtin_initialize builtinIncrementalElabs : IO.Ref NameSet ← IO.mkRef {}
+
+def addBuiltinIncrementalElab (decl : Name) : IO Unit := do
+  builtinIncrementalElabs.modify fun s => s.insert decl
+
+@[inherit_doc incrementalAttr, builtin_doc]
+builtin_initialize
+  registerBuiltinAttribute {
+    name            := `builtin_incremental
+    descr           := s!"(builtin) {incrementalAttr.attr.descr}"
+    applicationTime := .afterCompilation
+    add             := fun decl stx kind => do
+      Attribute.Builtin.ensureNoArgs stx
+      unless kind == AttributeKind.global do
+        throwError "invalid attribute 'builtin_incremental', must be global"
+      declareBuiltin decl <| mkApp (mkConst ``addBuiltinIncrementalElab) (toExpr decl)
+  }
+
+/-- Checks whether a declaration is annotated with `[builtin_incremental]` or `[incremental]`. -/
+def isIncrementalElab [Monad m] [MonadEnv m] [MonadLiftT IO m] (decl : Name) : m Bool :=
+  (return (← builtinIncrementalElabs.get (m := IO)).contains decl) <||>
+  (return incrementalAttr.hasTag (← getEnv) decl)
+
+export Term (TermElabM)
+
+builtin_initialize
+  registerTraceClass `Elab.implicitForall
+
+end Lean.Elab

backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Elab/Time.lean

@@ -0,0 +1,26 @@
+/-
+Copyright (c) 2021 Mario Carneiro. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Mario Carneiro
+-/
+prelude
+import Lean.Elab.Command
+
+/-!
+# Defines `#time` command.
+
+Time the elaboration of a command, and print the result (in milliseconds).
+-/
+
+namespace Lean.Elab.Time
+
+open Command
+
+@[builtin_command_elab Lean.Parser.timeCmd] def elabTimeCmd : CommandElab
+  | `(#time%$tk $stx:command) => do
+    let start ← IO.monoMsNow
+    elabCommand stx
+    logInfoAt tk m!"time: {(← IO.monoMsNow) - start}ms"
+  | _ => throwUnsupportedSyntax
+
+end Lean.Elab.Time

backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Elab/Util.lean

@@ -0,0 +1,249 @@
+/-
+Copyright (c) 2019 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import Lean.Parser.Command
+import Lean.KeyedDeclsAttribute
+import Lean.Elab.Exception
+
+namespace Lean
+
+def Syntax.prettyPrint (stx : Syntax) : Format :=
+  match stx.unsetTrailing.reprint with -- TODO use syntax pretty printer
+  | some str => format str.toFormat
+  | none     => format stx
+
+def MacroScopesView.format (view : MacroScopesView) (mainModule : Name) : Format :=
+  Std.format <|
+    if view.scopes.isEmpty then
+      view.name
+    else if view.mainModule == mainModule then
+      view.scopes.foldl Name.mkNum (view.name ++ view.imported)
+    else
+      view.scopes.foldl Name.mkNum (view.name ++ view.imported ++ view.mainModule)
+
+/--
+Two names are from the same lexical scope if their scoping information modulo `MacroScopesView.name`
+is equal.
+-/
+def MacroScopesView.equalScope (a b : MacroScopesView) : Bool :=
+  a.scopes == b.scopes && a.mainModule == b.mainModule && a.imported == b.imported
+
+namespace Elab
+
+def expandOptNamedPrio (stx : Syntax) : MacroM Nat :=
+  if stx.isNone then
+    return eval_prio default
+  else match stx[0] with
+    | `(Parser.Command.namedPrio| (priority := $prio)) => evalPrio prio
+    | _ => Macro.throwUnsupported
+
+structure MacroStackElem where
+  before : Syntax
+  after : Syntax
+
+abbrev MacroStack := List MacroStackElem
+
+/-- If `ref` does not have position information, then try to use macroStack -/
+def getBetterRef (ref : Syntax) (macroStack : MacroStack) : Syntax :=
+  match ref.getPos? with
+  | some _ => ref
+  | none   =>
+    match macroStack.find? (·.before.getPos? != none) with
+    | some elem => elem.before
+    | none      => ref
+
+register_builtin_option pp.macroStack : Bool := {
+  defValue := false
+  group    := "pp"
+  descr    := "display macro expansion stack"
+}
+
+def addMacroStack {m} [Monad m] [MonadOptions m] (msgData : MessageData) (macroStack : MacroStack) : m MessageData := do
+  if !pp.macroStack.get (← getOptions) then pure msgData else
+  match macroStack with
+  | []             => pure msgData
+  | stack@(top::_) =>
+    let msgData := msgData ++ Format.line ++ "with resulting expansion" ++ indentD top.after
+    pure $ stack.foldl
+      (fun (msgData : MessageData) (elem : MacroStackElem) =>
+        msgData ++ Format.line ++ "while expanding" ++ indentD elem.before)
+      msgData
+
+def checkSyntaxNodeKind [Monad m] [MonadEnv m] [MonadError m] (k : Name) : m Name := do
+  if Parser.isValidSyntaxNodeKind (← getEnv) k then pure k
+  else throwError "failed"
+
+def checkSyntaxNodeKindAtNamespaces [Monad m] [MonadEnv m] [MonadError m] (k : Name) : Name → m Name
+  | n@(.str p _) => checkSyntaxNodeKind (n ++ k) <|> checkSyntaxNodeKindAtNamespaces k p
+  | .anonymous   => checkSyntaxNodeKind k
+  | _ => throwError "failed"
+
+def checkSyntaxNodeKindAtCurrentNamespaces (k : Name) : AttrM Name := do
+  let ctx ← read
+  checkSyntaxNodeKindAtNamespaces k ctx.currNamespace
+
+def syntaxNodeKindOfAttrParam (defaultParserNamespace : Name) (stx : Syntax) : AttrM SyntaxNodeKind := do
+  let k ← Attribute.Builtin.getId stx
+  checkSyntaxNodeKindAtCurrentNamespaces k
+  <|>
+  checkSyntaxNodeKind (defaultParserNamespace ++ k)
+  <|>
+  throwError "invalid syntax node kind '{k}'"
+
+private unsafe def evalSyntaxConstantUnsafe (env : Environment) (opts : Options) (constName : Name) : ExceptT String Id Syntax :=
+  env.evalConstCheck Syntax opts `Lean.Syntax constName
+
+@[implemented_by evalSyntaxConstantUnsafe]
+opaque evalSyntaxConstant (env : Environment) (opts : Options) (constName : Name) : ExceptT String Id Syntax := throw ""
+
+unsafe def mkElabAttribute (γ) (attrBuiltinName attrName : Name) (parserNamespace : Name) (typeName : Name) (kind : String)
+    (attrDeclName : Name := by exact decl_name%) : IO (KeyedDeclsAttribute γ) :=
+  KeyedDeclsAttribute.init {
+    builtinName   := attrBuiltinName
+    name          := attrName
+    descr         := kind ++ " elaborator"
+    valueTypeName := typeName
+    evalKey       := fun _ stx => do
+      let kind ← syntaxNodeKindOfAttrParam parserNamespace stx
+      /- Recall that a `SyntaxNodeKind` is often the name of the parser, but this is not always true, and we must check it. -/
+      if (← getEnv).contains kind && (← getInfoState).enabled then
+        addConstInfo stx[1] kind none
+      return kind
+    onAdded       := fun builtin declName => do
+      if builtin then
+        declareBuiltinDocStringAndRanges declName
+  } attrDeclName
+
+unsafe def mkMacroAttributeUnsafe (ref : Name) : IO (KeyedDeclsAttribute Macro) :=
+  mkElabAttribute Macro `builtin_macro `macro Name.anonymous `Lean.Macro "macro" ref
+
+@[implemented_by mkMacroAttributeUnsafe]
+opaque mkMacroAttribute (ref : Name) : IO (KeyedDeclsAttribute Macro)
+
+/--
+Registers a macro expander for a given syntax node kind.
+
+A macro expander should have type `Lean.Macro` (which is `Lean.Syntax → Lean.MacroM Lean.Syntax`),
+i.e. should take syntax of the given syntax node kind as a parameter and produce different syntax
+in the same syntax category.
+
+The `macro_rules` and `macro` commands should usually be preferred over using this attribute
+directly.
+-/
+@[builtin_doc]
+builtin_initialize macroAttribute : KeyedDeclsAttribute Macro ← mkMacroAttribute decl_name%
+
+/--
+Try to expand macro at syntax tree root and return macro declaration name and new syntax if successful.
+Return none if all macros threw `Macro.Exception.unsupportedSyntax`.
+-/
+def expandMacroImpl? (env : Environment) : Syntax → MacroM (Option (Name × Except Macro.Exception Syntax)) := fun stx => do
+  for e in macroAttribute.getEntries env stx.getKind do
+    try
+      let stx' ← withFreshMacroScope (e.value stx)
+      return (e.declName, Except.ok stx')
+    catch
+      | Macro.Exception.unsupportedSyntax => pure ()
+      | ex                                => return (e.declName, Except.error ex)
+  return none
+
+class MonadMacroAdapter (m : Type → Type) where
+  getCurrMacroScope                  : m MacroScope
+  getNextMacroScope                  : m MacroScope
+  setNextMacroScope                  : MacroScope → m Unit
+
+@[always_inline]
+instance (m n) [MonadLift m n] [MonadMacroAdapter m] : MonadMacroAdapter n := {
+  getCurrMacroScope := liftM (MonadMacroAdapter.getCurrMacroScope : m _)
+  getNextMacroScope := liftM (MonadMacroAdapter.getNextMacroScope : m _)
+  setNextMacroScope := fun s => liftM (MonadMacroAdapter.setNextMacroScope s : m _)
+}
+
+def liftMacroM [Monad m] [MonadMacroAdapter m] [MonadEnv m] [MonadRecDepth m] [MonadError m] [MonadResolveName m] [MonadTrace m] [MonadOptions m] [AddMessageContext m] [MonadLiftT IO m] (x : MacroM α) : m α := do
+  let env  ← getEnv
+  let currNamespace ← getCurrNamespace
+  let openDecls ← getOpenDecls
+  let methods := Macro.mkMethods {
+    -- TODO: record recursive expansions in info tree?
+    expandMacro?     := fun stx => do
+      match (← expandMacroImpl? env stx) with
+      | some (_, stx?) => liftExcept stx?
+      | none           => return none
+    hasDecl          := fun declName => return env.contains declName
+    getCurrNamespace := return currNamespace
+    resolveNamespace := fun n => return ResolveName.resolveNamespace env currNamespace openDecls n
+    resolveGlobalName := fun n => return ResolveName.resolveGlobalName env currNamespace openDecls n
+  }
+  match x { methods        := methods
+            ref            := ← getRef
+            currMacroScope := ← MonadMacroAdapter.getCurrMacroScope
+            mainModule     := env.mainModule
+            currRecDepth   := ← MonadRecDepth.getRecDepth
+            maxRecDepth    := ← MonadRecDepth.getMaxRecDepth
+          } { macroScope := (← MonadMacroAdapter.getNextMacroScope) } with
+  | EStateM.Result.error Macro.Exception.unsupportedSyntax _ => throwUnsupportedSyntax
+  | EStateM.Result.error (Macro.Exception.error ref msg) _   =>
+    if msg == maxRecDepthErrorMessage then
+      -- Make sure we can detect exception using `Exception.isMaxRecDepth`
+      throwMaxRecDepthAt ref
+    else
+      throwErrorAt ref msg
+  | EStateM.Result.ok a s =>
+    MonadMacroAdapter.setNextMacroScope s.macroScope
+    s.traceMsgs.reverse.forM fun (clsName, msg) => trace clsName fun _ => msg
+    return a
+
+@[inline] def adaptMacro {m : Type → Type} [Monad m] [MonadMacroAdapter m] [MonadEnv m] [MonadRecDepth m] [MonadError m]  [MonadResolveName m] [MonadTrace m] [MonadOptions m] [AddMessageContext m] [MonadLiftT IO m] (x : Macro) (stx : Syntax) : m Syntax :=
+  liftMacroM (x stx)
+
+partial def mkUnusedBaseName (baseName : Name) : MacroM Name := do
+  let currNamespace ← Macro.getCurrNamespace
+  if ← Macro.hasDecl (currNamespace ++ baseName) then
+    let rec loop (idx : Nat) := do
+      let name := baseName.appendIndexAfter idx
+      if ← Macro.hasDecl (currNamespace ++ name) then
+        loop (idx+1)
+      else
+        return name
+    loop 1
+  else
+    return baseName
+
+def logException [Monad m] [MonadLog m] [AddMessageContext m] [MonadOptions m] [MonadLiftT IO m] (ex : Exception) : m Unit := do
+  match ex with
+  | Exception.error ref msg => logErrorAt ref msg
+  | Exception.internal id _ =>
+    unless isAbortExceptionId id || ex.isInterrupt do
+      let name ← id.getName
+      logError m!"internal exception: {name}"
+
+/--
+If `x` throws an exception, catch it and turn it into a log message (using `logException`).
+-/
+def withLogging [Monad m] [MonadLog m] [MonadExcept Exception m] [AddMessageContext m] [MonadOptions m] [MonadLiftT IO m]
+    (x : m Unit) : m Unit := do
+  try x catch ex => logException ex
+
+def nestedExceptionToMessageData [Monad m] [MonadLog m] (ex : Exception) : m MessageData := do
+  let pos ← getRefPos
+  match ex.getRef.getPos? with
+  | none       => return ex.toMessageData
+  | some exPos =>
+    if pos == exPos then
+      return ex.toMessageData
+    else
+      let exPosition := (← getFileMap).toPosition exPos
+      return m!"{exPosition.line}:{exPosition.column} {ex.toMessageData}"
+
+def throwErrorWithNestedErrors [MonadError m] [Monad m] [MonadLog m] (msg : MessageData) (exs : Array Exception) : m α := do
+  throwError "{msg}, errors {toMessageList (← exs.mapM fun | ex => nestedExceptionToMessageData ex)}"
+
+builtin_initialize
+  registerTraceClass `Elab
+  registerTraceClass `Elab.step
+  registerTraceClass `Elab.step.result (inherited := true)
+
+end Lean.Elab

backend/core/leanprover--lean4---v4.22.0/src/lean/Lean/Elab/WhereFinally.lean

@@ -0,0 +1,29 @@
+/-
+Copyright (c) 2025 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Sebastian Graf
+-/
+prelude
+import Lean.Parser.Term
+
+namespace Lean.Elab
+
+structure WhereFinallyView where
+  ref       : Syntax
+  tactic    : TSyntax ``Lean.Parser.Tactic.tacticSeq
+  deriving Inhabited
+
+def WhereFinallyView.none : WhereFinallyView := { ref := .missing, tactic := ⟨.missing⟩ }
+
+def WhereFinallyView.isNone (o : WhereFinallyView) : Bool := o.ref.isMissing && o.tactic.raw.isMissing
+
+/-- Creates a view of the `finally` section of a `whereDecls` syntax object -/
+def mkWhereFinallyView {m} [Monad m] [MonadError m] (stx : TSyntax ``Parser.Term.whereDecls) : m WhereFinallyView := do
+  -- Fail gracefully upon partial parses/missing where or finally sections
+  let whereFinally := stx.raw[2][0]
+  if whereFinally.isMissing then
+    return { ref := stx, tactic := ⟨.missing⟩ }
+  if !whereFinally[2][0].isMissing then
+    throwErrorAt stx "`where ... finally` does not currently support any named sub-sections `| sectionName => ...`"
+  let tactic := ⟨whereFinally[1]⟩
+  return { ref := whereFinally, tactic }

external/alphageometry/.venv-ag/Lib/site-packages/absl/app.py

@@ -0,0 +1,488 @@
+# Copyright 2017 The Abseil Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Generic entry point for Abseil Python applications.
+
+To use this module, define a ``main`` function with a single ``argv`` argument
+and call ``app.run(main)``. For example::
+
+    def main(argv):
+      if len(argv) > 1:
+        raise app.UsageError('Too many command-line arguments.')
+
+    if __name__ == '__main__':
+      app.run(main)
+"""
+
+import collections
+import errno
+import os
+import pdb
+import sys
+import textwrap
+import traceback
+
+from absl import command_name
+from absl import flags
+from absl import logging
+
+try:
+  import faulthandler
+except ImportError:
+  faulthandler = None
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_boolean('run_with_pdb', False, 'Set to true for PDB debug mode')
+flags.DEFINE_boolean('pdb_post_mortem', False,
+                     'Set to true to handle uncaught exceptions with PDB '
+                     'post mortem.')
+flags.DEFINE_alias('pdb', 'pdb_post_mortem')
+flags.DEFINE_boolean('run_with_profiling', False,
+                     'Set to true for profiling the script. '
+                     'Execution will be slower, and the output format might '
+                     'change over time.')
+flags.DEFINE_string('profile_file', None,
+                    'Dump profile information to a file (for python -m '
+                    'pstats). Implies --run_with_profiling.')
+flags.DEFINE_boolean('use_cprofile_for_profiling', True,
+                     'Use cProfile instead of the profile module for '
+                     'profiling. This has no effect unless '
+                     '--run_with_profiling is set.')
+flags.DEFINE_boolean('only_check_args', False,
+                     'Set to true to validate args and exit.',
+                     allow_hide_cpp=True)
+
+
+# If main() exits via an abnormal exception, call into these
+# handlers before exiting.
+EXCEPTION_HANDLERS = []
+
+
+class Error(Exception):
+  pass
+
+
+class UsageError(Error):
+  """Exception raised when the arguments supplied by the user are invalid.
+
+  Raise this when the arguments supplied are invalid from the point of
+  view of the application. For example when two mutually exclusive
+  flags have been supplied or when there are not enough non-flag
+  arguments. It is distinct from flags.Error which covers the lower
+  level of parsing and validating individual flags.
+  """
+
+  def __init__(self, message, exitcode=1):
+    super().__init__(message)
+    self.exitcode = exitcode
+
+
+class HelpFlag(flags.BooleanFlag):
+  """Special boolean flag that displays usage and raises SystemExit."""
+  NAME = 'help'
+  SHORT_NAME = '?'
+
+  def __init__(self):
+    super().__init__(
+        self.NAME,
+        False,
+        'show this help',
+        short_name=self.SHORT_NAME,
+        allow_hide_cpp=True,
+    )
+
+  def parse(self, arg):
+    if self._parse(arg):
+      usage(shorthelp=True, writeto_stdout=True)
+      # Advertise --helpfull on stdout, since usage() was on stdout.
+      print()
+      print('Try --helpfull to get a list of all flags.')
+      sys.exit(1)
+
+
+class HelpshortFlag(HelpFlag):
+  """--helpshort is an alias for --help."""
+  NAME = 'helpshort'
+  SHORT_NAME = None
+
+
+class HelpfullFlag(flags.BooleanFlag):
+  """Display help for flags in the main module and all dependent modules."""
+
+  def __init__(self):
+    super().__init__('helpfull', False, 'show full help', allow_hide_cpp=True)
+
+  def parse(self, arg):
+    if self._parse(arg):
+      usage(writeto_stdout=True)
+      sys.exit(1)
+
+
+class HelpXMLFlag(flags.BooleanFlag):
+  """Similar to HelpfullFlag, but generates output in XML format."""
+
+  def __init__(self):
+    super().__init__(
+        'helpxml',
+        False,
+        'like --helpfull, but generates XML output',
+        allow_hide_cpp=True,
+    )
+
+  def parse(self, arg):
+    if self._parse(arg):
+      flags.FLAGS.write_help_in_xml_format(sys.stdout)
+      sys.exit(1)
+
+
+def parse_flags_with_usage(args):
+  """Tries to parse the flags, print usage, and exit if unparsable.
+
+  Args:
+    args: [str], a non-empty list of the command line arguments including
+        program name.
+
+  Returns:
+    [str], a non-empty list of remaining command line arguments after parsing
+    flags, including program name.
+  """
+  try:
+    return FLAGS(args)
+  except flags.Error as error:
+    message = str(error)
+    if '\n' in message:
+      final_message = 'FATAL Flags parsing error:\n%s\n' % textwrap.indent(
+          message, '  ')
+    else:
+      final_message = 'FATAL Flags parsing error: %s\n' % message
+    sys.stderr.write(final_message)
+    sys.stderr.write('Pass --helpshort or --helpfull to see help on flags.\n')
+    sys.exit(1)
+
+
+_define_help_flags_called = False
+
+
+def define_help_flags():
+  """Registers help flags. Idempotent."""
+  # Use a global to ensure idempotence.
+  global _define_help_flags_called
+
+  if not _define_help_flags_called:
+    flags.DEFINE_flag(HelpFlag())
+    flags.DEFINE_flag(HelpshortFlag())  # alias for --help
+    flags.DEFINE_flag(HelpfullFlag())
+    flags.DEFINE_flag(HelpXMLFlag())
+    _define_help_flags_called = True
+
+
+def _register_and_parse_flags_with_usage(
+    argv=None,
+    flags_parser=parse_flags_with_usage,
+):
+  """Registers help flags, parses arguments and shows usage if appropriate.
+
+  This also calls sys.exit(0) if flag --only_check_args is True.
+
+  Args:
+    argv: [str], a non-empty list of the command line arguments including
+      program name, sys.argv is used if None.
+    flags_parser: Callable[[List[str]], Any], the function used to parse flags.
+      The return value of this function is passed to `main` untouched. It must
+      guarantee FLAGS is parsed after this function is called.
+
+  Returns:
+    The return value of `flags_parser`. When using the default `flags_parser`,
+    it returns the following:
+    [str], a non-empty list of remaining command line arguments after parsing
+    flags, including program name.
+
+  Raises:
+    Error: Raised when flags_parser is called, but FLAGS is not parsed.
+    SystemError: Raised when it's called more than once.
+  """
+  # fmt: on
+  if _register_and_parse_flags_with_usage.done:
+    raise SystemError('Flag registration can be done only once.')
+
+  define_help_flags()
+
+  original_argv = sys.argv if argv is None else argv
+  args_to_main = flags_parser(original_argv)
+  if not FLAGS.is_parsed():
+    raise Error('FLAGS must be parsed after flags_parser is called.')
+
+  # Exit when told so.
+  if FLAGS.only_check_args:
+    sys.exit(0)
+  # Immediately after flags are parsed, bump verbosity to INFO if the flag has
+  # not been set.
+  if FLAGS['verbosity'].using_default_value:
+    FLAGS.verbosity = 0
+  _register_and_parse_flags_with_usage.done = True
+
+  return args_to_main
+
+_register_and_parse_flags_with_usage.done = False
+
+
+def _run_main(main, argv):
+  """Calls main, optionally with pdb or profiler."""
+  if FLAGS.run_with_pdb:
+    sys.exit(pdb.runcall(main, argv))
+  elif FLAGS.run_with_profiling or FLAGS.profile_file:
+    # Avoid import overhead since most apps (including performance-sensitive
+    # ones) won't be run with profiling.
+    # pylint: disable=g-import-not-at-top
+    import atexit
+    if FLAGS.use_cprofile_for_profiling:
+      import cProfile as profile
+    else:
+      import profile
+    profiler = profile.Profile()
+    if FLAGS.profile_file:
+      atexit.register(profiler.dump_stats, FLAGS.profile_file)
+    else:
+      atexit.register(profiler.print_stats)
+    sys.exit(profiler.runcall(main, argv))
+  else:
+    sys.exit(main(argv))
+
+
+def _call_exception_handlers(exception):
+  """Calls any installed exception handlers."""
+  for handler in EXCEPTION_HANDLERS:
+    try:
+      if handler.wants(exception):
+        handler.handle(exception)
+    except:  # pylint: disable=bare-except
+      try:
+        # We don't want to stop for exceptions in the exception handlers but
+        # we shouldn't hide them either.
+        logging.error(traceback.format_exc())
+      except:  # pylint: disable=bare-except
+        # In case even the logging statement fails, ignore.
+        pass
+
+
+def run(
+    main,
+    argv=None,
+    flags_parser=parse_flags_with_usage,
+):
+  """Begins executing the program.
+
+  Args:
+    main: The main function to execute. It takes an single argument "argv",
+        which is a list of command line arguments with parsed flags removed.
+        The return value is passed to `sys.exit`, and so for example
+        a return value of 0 or None results in a successful termination, whereas
+        a return value of 1 results in abnormal termination.
+        For more details, see https://docs.python.org/3/library/sys#sys.exit
+    argv: A non-empty list of the command line arguments including program name,
+        sys.argv is used if None.
+    flags_parser: Callable[[List[str]], Any], the function used to parse flags.
+        The return value of this function is passed to `main` untouched.
+        It must guarantee FLAGS is parsed after this function is called.
+        Should be passed as a keyword-only arg which will become mandatory in a
+        future release.
+  - Parses command line flags with the flag module.
+  - If there are any errors, prints usage().
+  - Calls main() with the remaining arguments.
+  - If main() raises a UsageError, prints usage and the error message.
+  """
+  # fmt: on
+  try:
+    args = _run_init(
+        sys.argv if argv is None else argv,
+        flags_parser,
+    )
+    while _init_callbacks:
+      callback = _init_callbacks.popleft()
+      callback()
+    try:
+      _run_main(main, args)
+    except UsageError as error:
+      usage(shorthelp=True, detailed_error=error, exitcode=error.exitcode)
+    except:
+      exc = sys.exc_info()[1]
+      # Don't try to post-mortem debug successful SystemExits, since those
+      # mean there wasn't actually an error. In particular, the test framework
+      # raises SystemExit(False) even if all tests passed.
+      if isinstance(exc, SystemExit) and not exc.code:
+        raise
+
+      # Check the tty so that we don't hang waiting for input in an
+      # non-interactive scenario.
+      if FLAGS.pdb_post_mortem and sys.stdout.isatty():
+        traceback.print_exc()
+        print()
+        print(' *** Entering post-mortem debugging ***')
+        print()
+        pdb.post_mortem()
+      raise
+  except Exception as e:
+    _call_exception_handlers(e)
+    raise
+
+# Callbacks which have been deferred until after _run_init has been called.
+_init_callbacks = collections.deque()
+
+
+def call_after_init(callback):
+  """Calls the given callback only once ABSL has finished initialization.
+
+  If ABSL has already finished initialization when ``call_after_init`` is
+  called then the callback is executed immediately, otherwise `callback` is
+  stored to be executed after ``app.run`` has finished initializing (aka. just
+  before the main function is called).
+
+  If called after ``app.run``, this is equivalent to calling ``callback()`` in
+  the caller thread. If called before ``app.run``, callbacks are run
+  sequentially (in an undefined order) in the same thread as ``app.run``.
+
+  Args:
+    callback: a callable to be called once ABSL has finished initialization.
+      This may be immediate if initialization has already finished. It
+      takes no arguments and returns nothing.
+  """
+  if _run_init.done:
+    callback()
+  else:
+    _init_callbacks.append(callback)
+
+
+def _run_init(
+    argv,
+    flags_parser,
+):
+  """Does one-time initialization and re-parses flags on rerun."""
+  if _run_init.done:
+    return flags_parser(argv)
+  command_name.make_process_name_useful()
+  # Set up absl logging handler.
+  logging.use_absl_handler()
+  args = _register_and_parse_flags_with_usage(
+      argv=argv,
+      flags_parser=flags_parser,
+  )
+  if faulthandler:
+    try:
+      faulthandler.enable()
+    except Exception:  # pylint: disable=broad-except
+      # Some tests verify stderr output very closely, so don't print anything.
+      # Disabled faulthandler is a low-impact error.
+      pass
+  _run_init.done = True
+  return args
+
+
+_run_init.done = False
+
+
+def usage(shorthelp=False, writeto_stdout=False, detailed_error=None,
+          exitcode=None):
+  """Writes __main__'s docstring to stderr with some help text.
+
+  Args:
+    shorthelp: bool, if True, prints only flags from the main module,
+        rather than all flags.
+    writeto_stdout: bool, if True, writes help message to stdout,
+        rather than to stderr.
+    detailed_error: str, additional detail about why usage info was presented.
+    exitcode: optional integer, if set, exits with this status code after
+        writing help.
+  """
+  if writeto_stdout:
+    stdfile = sys.stdout
+  else:
+    stdfile = sys.stderr
+
+  doc = sys.modules['__main__'].__doc__
+  if not doc:
+    doc = '\nUSAGE: %s [flags]\n' % sys.argv[0]
+    doc = flags.text_wrap(doc, indent='       ', firstline_indent='')
+  else:
+    # Replace all '%s' with sys.argv[0], and all '%%' with '%'.
+    num_specifiers = doc.count('%') - 2 * doc.count('%%')
+    try:
+      doc %= (sys.argv[0],) * num_specifiers
+    except (OverflowError, TypeError, ValueError):
+      # Just display the docstring as-is.
+      pass
+  if shorthelp:
+    flag_str = FLAGS.main_module_help()
+  else:
+    flag_str = FLAGS.get_help()
+  try:
+    stdfile.write(doc)
+    if flag_str:
+      stdfile.write('\nflags:\n')
+      stdfile.write(flag_str)
+    stdfile.write('\n')
+    if detailed_error is not None:
+      stdfile.write('\n%s\n' % detailed_error)
+  except OSError as e:
+    # We avoid printing a huge backtrace if we get EPIPE, because
+    # "foo.par --help | less" is a frequent use case.
+    if e.errno != errno.EPIPE:
+      raise
+  if exitcode is not None:
+    sys.exit(exitcode)
+
+
+class ExceptionHandler:
+  """Base exception handler from which other may inherit."""
+
+  def wants(self, exc):
+    """Returns whether this handler wants to handle the exception or not.
+
+    This base class returns True for all exceptions by default. Override in
+    subclass if it wants to be more selective.
+
+    Args:
+      exc: Exception, the current exception.
+    """
+    del exc  # Unused.
+    return True
+
+  def handle(self, exc):
+    """Do something with the current exception.
+
+    Args:
+      exc: Exception, the current exception
+
+    This method must be overridden.
+    """
+    raise NotImplementedError()
+
+
+def install_exception_handler(handler):
+  """Installs an exception handler.
+
+  Args:
+    handler: ExceptionHandler, the exception handler to install.
+
+  Raises:
+    TypeError: Raised when the handler was not of the correct type.
+
+  All installed exception handlers will be called if main() exits via
+  an abnormal exception, i.e. not one of SystemExit, KeyboardInterrupt,
+  FlagsError or UsageError.
+  """
+  if not isinstance(handler, ExceptionHandler):
+    raise TypeError('handler of type %s does not inherit from ExceptionHandler'
+                    % type(handler))
+  EXCEPTION_HANDLERS.append(handler)

external/alphageometry/.venv-ag/Lib/site-packages/absl/app.pyi

@@ -0,0 +1,88 @@
+from typing import Any, Callable, Collection, Iterable, List, NoReturn, Optional, TypeVar, Union, overload
+
+from absl.flags import _flag
+
+_MainArgs = TypeVar('_MainArgs')
+_Exc = TypeVar('_Exc', bound=Exception)
+
+class ExceptionHandler():
+
+  def wants(self, exc: _Exc) -> bool:
+    ...
+
+  def handle(self, exc: _Exc):
+    ...
+
+EXCEPTION_HANDLERS: List[ExceptionHandler] = ...
+
+class HelpFlag(_flag.BooleanFlag):
+  def __init__(self):
+    ...
+
+class HelpshortFlag(HelpFlag):
+  ...
+
+class HelpfullFlag(_flag.BooleanFlag):
+  def __init__(self):
+    ...
+
+class HelpXMLFlag(_flag.BooleanFlag):
+  def __init__(self):
+    ...
+
+def define_help_flags() -> None:
+  ...
+
+@overload
+def usage(shorthelp: Union[bool, int] = ...,
+          writeto_stdout: Union[bool, int] = ...,
+          detailed_error: Optional[Any] = ...,
+          exitcode: None = ...) -> None:
+  ...
+
+@overload
+def usage(shorthelp: Union[bool, int],
+          writeto_stdout: Union[bool, int],
+          detailed_error: Optional[Any],
+          exitcode: int) -> NoReturn:
+  ...
+
+@overload
+def usage(shorthelp: Union[bool, int] = ...,
+          writeto_stdout: Union[bool, int] = ...,
+          detailed_error: Optional[Any] = ...,
+          *,
+          exitcode: int) -> NoReturn:
+  ...
+
+def install_exception_handler(handler: ExceptionHandler) -> None:
+  ...
+
+class Error(Exception):
+  ...
+
+class UsageError(Error):
+  exitcode: int
+
+def parse_flags_with_usage(args: List[str]) -> List[str]:
+  ...
+
+def call_after_init(callback: Callable[[], Any]) -> None:
+  ...
+
+# Without the flag_parser argument, `main` should require a List[str].
+@overload
+def run(
+    main: Callable[[List[str]], Any],
+    argv: Optional[List[str]] = ...,
+) -> NoReturn:
+  ...
+
+@overload
+def run(
+    main: Callable[[_MainArgs], Any],
+    argv: Optional[List[str]] = ...,
+    *,
+    flags_parser: Callable[[List[str]], _MainArgs],
+) -> NoReturn:
+  ...

external/alphageometry/.venv-ag/Lib/site-packages/absl/command_name.py

@@ -0,0 +1,63 @@
+# Copyright 2017 The Abseil Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A tiny stand alone library to change the kernel process name on Linux."""
+
+import os
+import sys
+
+# This library must be kept small and stand alone.  It is used by small things
+# that require no extension modules.
+
+
+def make_process_name_useful():
+  """Sets the process name to something better than 'python' if possible."""
+  set_kernel_process_name(os.path.basename(sys.argv[0]))
+
+
+def set_kernel_process_name(name):
+  """Changes the Kernel's /proc/self/status process name on Linux.
+
+  The kernel name is NOT what will be shown by the ps or top command.
+  It is a 15 character string stored in the kernel's process table that
+  is included in the kernel log when a process is OOM killed.
+  The first 15 bytes of name are used.  Non-ASCII unicode is replaced with '?'.
+
+  Does nothing if /proc/self/comm cannot be written or prctl() fails.
+
+  Args:
+    name: bytes|unicode, the Linux kernel's command name to set.
+  """
+  if not isinstance(name, bytes):
+    name = name.encode('ascii', 'replace')
+  try:
+    # This is preferred to using ctypes to try and call prctl() when possible.
+    with open('/proc/self/comm', 'wb') as proc_comm:
+      proc_comm.write(name[:15])
+  except OSError:
+    try:
+      import ctypes  # pylint: disable=g-import-not-at-top
+    except ImportError:
+      return  # No ctypes.
+    try:
+      libc = ctypes.CDLL('libc.so.6')
+    except OSError:
+      return  # No libc.so.6.
+    pr_set_name = ctypes.c_ulong(15)  # linux/prctl.h PR_SET_NAME value.
+    zero = ctypes.c_ulong(0)
+    try:
+      libc.prctl(pr_set_name, name, zero, zero, zero)
+      # Ignore the prctl return value.  Nothing we can do if it errored.
+    except AttributeError:
+      return  # No prctl.

external/alphageometry/.venv-ag/Lib/site-packages/distutils-precedence.pth

@@ -0,0 +1 @@
+import os; var = 'SETUPTOOLS_USE_DISTUTILS'; enabled = os.environ.get(var, 'local') == 'local'; enabled and __import__('_distutils_hack').add_shim(); 

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2018 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
+from jax.experimental.x64_context import (
+  enable_x64 as enable_x64,
+  disable_x64 as disable_x64,
+)
+from jax._src.callback import (
+  io_callback as io_callback
+)
+from jax._src.earray import (
+    EArray as EArray
+)

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_api/__init__.py

@@ -0,0 +1,213 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This module includes experimental JAX support for the `Python array API standard`_.
+Support for this is currently experimental and not fully complete.
+
+Example Usage::
+
+  >>> from jax.experimental import array_api as xp
+
+  >>> xp.__array_api_version__
+  '2023.12'
+
+  >>> arr = xp.arange(1000)
+
+  >>> arr.sum()
+  Array(499500, dtype=int32)
+
+The ``xp`` namespace is the array API compliant analog of :mod:`jax.numpy`,
+and implements most of the API listed in the standard.
+
+.. _Python array API standard: https://data-apis.org/array-api/latest/
+"""
+
+from __future__ import annotations
+
+from jax.experimental.array_api._version import __array_api_version__ as __array_api_version__
+
+from jax.experimental.array_api import fft as fft
+from jax.experimental.array_api import linalg as linalg
+
+from jax.numpy import (
+    abs as abs,
+    acos as acos,
+    acosh as acosh,
+    add as add,
+    all as all,
+    any as any,
+    argmax as argmax,
+    argmin as argmin,
+    argsort as argsort,
+    asin as asin,
+    asinh as asinh,
+    atan as atan,
+    atan2 as atan2,
+    atanh as atanh,
+    bitwise_and as bitwise_and,
+    bitwise_invert as bitwise_invert,
+    bitwise_left_shift as bitwise_left_shift,
+    bitwise_or as bitwise_or,
+    bitwise_right_shift as bitwise_right_shift,
+    bitwise_xor as bitwise_xor,
+    bool as bool,
+    broadcast_arrays as broadcast_arrays,
+    broadcast_to as broadcast_to,
+    can_cast as can_cast,
+    complex128 as complex128,
+    complex64 as complex64,
+    concat as concat,
+    conj as conj,
+    copysign as copysign,
+    cos as cos,
+    cosh as cosh,
+    cumulative_sum as cumulative_sum,
+    divide as divide,
+    e as e,
+    empty as empty,
+    empty_like as empty_like,
+    equal as equal,
+    exp as exp,
+    expand_dims as expand_dims,
+    expm1 as expm1,
+    flip as flip,
+    float32 as float32,
+    float64 as float64,
+    floor_divide as floor_divide,
+    from_dlpack as from_dlpack,
+    full as full,
+    full_like as full_like,
+    greater as greater,
+    greater_equal as greater_equal,
+    iinfo as iinfo,
+    imag as imag,
+    inf as inf,
+    int16 as int16,
+    int32 as int32,
+    int64 as int64,
+    int8 as int8,
+    isdtype as isdtype,
+    isfinite as isfinite,
+    isinf as isinf,
+    isnan as isnan,
+    less as less,
+    less_equal as less_equal,
+    log as log,
+    log10 as log10,
+    log1p as log1p,
+    log2 as log2,
+    logaddexp as logaddexp,
+    logical_and as logical_and,
+    logical_not as logical_not,
+    logical_or as logical_or,
+    logical_xor as logical_xor,
+    matmul as matmul,
+    matrix_transpose as matrix_transpose,
+    max as max,
+    maximum as maximum,
+    mean as mean,
+    meshgrid as meshgrid,
+    min as min,
+    minimum as minimum,
+    moveaxis as moveaxis,
+    multiply as multiply,
+    nan as nan,
+    negative as negative,
+    newaxis as newaxis,
+    nonzero as nonzero,
+    not_equal as not_equal,
+    ones as ones,
+    ones_like as ones_like,
+    permute_dims as permute_dims,
+    pi as pi,
+    positive as positive,
+    pow as pow,
+    prod as prod,
+    real as real,
+    remainder as remainder,
+    repeat as repeat,
+    result_type as result_type,
+    roll as roll,
+    round as round,
+    searchsorted as searchsorted,
+    sign as sign,
+    signbit as signbit,
+    sin as sin,
+    sinh as sinh,
+    sort as sort,
+    sqrt as sqrt,
+    square as square,
+    squeeze as squeeze,
+    stack as stack,
+    subtract as subtract,
+    sum as sum,
+    take as take,
+    tan as tan,
+    tanh as tanh,
+    tensordot as tensordot,
+    tile as tile,
+    tril as tril,
+    triu as triu,
+    uint16 as uint16,
+    uint32 as uint32,
+    uint64 as uint64,
+    uint8 as uint8,
+    unique_all as unique_all,
+    unique_counts as unique_counts,
+    unique_inverse as unique_inverse,
+    unique_values as unique_values,
+    unstack as unstack,
+    vecdot as vecdot,
+    where as where,
+    zeros as zeros,
+    zeros_like as zeros_like,
+)
+
+from jax.experimental.array_api._manipulation_functions import (
+    reshape as reshape,
+)
+
+from jax.experimental.array_api._creation_functions import (
+    arange as arange,
+    asarray as asarray,
+    eye as eye,
+    linspace as linspace,
+)
+
+from jax.experimental.array_api._data_type_functions import (
+    astype as astype,
+    finfo as finfo,
+)
+
+from jax.experimental.array_api._elementwise_functions import (
+    ceil as ceil,
+    clip as clip,
+    floor as floor,
+    hypot as hypot,
+    trunc as trunc,
+)
+
+from jax.experimental.array_api._statistical_functions import (
+    std as std,
+    var as var,
+)
+
+from jax.experimental.array_api._utility_functions import (
+    __array_namespace_info__ as __array_namespace_info__,
+)
+
+from jax.experimental.array_api import _array_methods
+_array_methods.add_array_object_methods()
+del _array_methods

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_api/_array_methods.py

@@ -0,0 +1,45 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Any
+
+import jax
+from jax._src.array import ArrayImpl
+from jax.experimental.array_api._version import __array_api_version__
+from jax.sharding import Sharding
+
+from jax._src.lib import xla_extension as xe
+
+
+def _array_namespace(self, /, *, api_version: None | str = None):
+  if api_version is not None and api_version != __array_api_version__:
+    raise ValueError(f"{api_version=!r} is not available; "
+                     f"available versions are: {[__array_api_version__]}")
+  return jax.experimental.array_api
+
+
+def _to_device(self, device: xe.Device | Sharding | None, *,
+               stream: int | Any | None = None):
+  if stream is not None:
+    raise NotImplementedError("stream argument of array.to_device()")
+  return jax.device_put(self, device)
+
+
+def add_array_object_methods():
+  # TODO(jakevdp): set on tracers as well?
+  setattr(ArrayImpl, "__array_namespace__", _array_namespace)
+  setattr(ArrayImpl, "to_device", _to_device)
+  setattr(ArrayImpl, "device", property(lambda self: self.sharding))

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_api/_creation_functions.py

@@ -0,0 +1,31 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+
+# TODO(micky774): Deprecate after adding device argument to jax.numpy functions
+def arange(start, /, stop=None, step=1, *, dtype=None, device=None):
+  return jax.device_put(jnp.arange(start, stop, step, dtype=dtype), device=device)
+
+def asarray(obj, /, *, dtype=None, device=None, copy=None):
+  return jax.device_put(jnp.array(obj, dtype=dtype, copy=copy), device=device)
+
+def eye(n_rows, n_cols=None, /, *, k=0, dtype=None, device=None):
+  return jax.device_put(jnp.eye(n_rows, n_cols, k=k, dtype=dtype), device=device)
+
+def linspace(start, stop, /, num, *, dtype=None, device=None, endpoint=True):
+  return jax.device_put(jnp.linspace(start, stop, num=num, dtype=dtype, endpoint=endpoint), device=device)

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_api/_data_type_functions.py

@@ -0,0 +1,78 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import builtins
+from typing import NamedTuple
+import numpy as np
+
+import jax.numpy as jnp
+
+from jax._src.lib import xla_client as xc
+from jax._src.sharding import Sharding
+from jax._src import dtypes as _dtypes
+
+# TODO(micky774): Update jax.numpy dtypes to dtype *objects*
+bool = np.dtype('bool')
+int8 = np.dtype('int8')
+int16 = np.dtype('int16')
+int32 = np.dtype('int32')
+int64 = np.dtype('int64')
+uint8 = np.dtype('uint8')
+uint16 = np.dtype('uint16')
+uint32 = np.dtype('uint32')
+uint64 = np.dtype('uint64')
+float32 = np.dtype('float32')
+float64 = np.dtype('float64')
+complex64 = np.dtype('complex64')
+complex128 = np.dtype('complex128')
+
+
+# TODO(micky774): Remove when jax.numpy.astype is deprecation is completed
+def astype(x, dtype, /, *, copy: builtins.bool = True, device: xc.Device | Sharding | None = None):
+  src_dtype = x.dtype if hasattr(x, "dtype") else _dtypes.dtype(x)
+  if (
+    src_dtype is not None
+    and _dtypes.isdtype(src_dtype, "complex floating")
+    and _dtypes.isdtype(dtype, ("integral", "real floating"))
+  ):
+    raise ValueError(
+      "Casting from complex to non-complex dtypes is not permitted. Please "
+      "first use jnp.real or jnp.imag to take the real/imaginary component of "
+      "your input."
+    )
+  return jnp.astype(x, dtype, copy=copy, device=device)
+
+
+class FInfo(NamedTuple):
+    bits: int
+    eps: float
+    max: float
+    min: float
+    smallest_normal: float
+    dtype: jnp.dtype
+
+# TODO(micky774): Update jax.numpy.finfo so that its attributes are python
+# floats
+def finfo(type, /) -> FInfo:
+  info = jnp.finfo(type)
+  return FInfo(
+    bits=info.bits,
+    eps=float(info.eps),
+    max=float(info.max),
+    min=float(info.min),
+    smallest_normal=float(info.smallest_normal),
+    dtype=jnp.dtype(type)
+  )

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_api/_elementwise_functions.py

@@ -0,0 +1,75 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import jax
+from jax.numpy import isdtype
+from jax._src.dtypes import issubdtype
+from jax._src.numpy.util import promote_args
+
+
+# TODO(micky774): Update jnp.ceil to preserve integral dtype
+def ceil(x, /):
+  """Rounds each element x_i of the input array x to the smallest (i.e., closest to -infinity) integer-valued number that is not less than x_i."""
+  x, = promote_args("ceil", x)
+  if isdtype(x.dtype, "integral"):
+    return x
+  return jax.numpy.ceil(x)
+
+
+# TODO(micky774): Remove when jnp.clip deprecation is completed
+# (began 2024-4-2) and default behavior is Array API 2023 compliant
+def clip(x, /, min=None, max=None):
+  """Returns the complex conjugate for each element x_i of the input array x."""
+  x, = promote_args("clip", x)
+
+  if any(jax.numpy.iscomplexobj(t) for t in (x, min, max)):
+    raise ValueError(
+      "Clip received a complex value either through the input or the min/max "
+      "keywords. Complex values have no ordering and cannot be clipped. "
+      "Please convert to a real value or array by taking the real or "
+      "imaginary components via jax.numpy.real/imag respectively."
+    )
+  return jax.numpy.clip(x, min=min, max=max)
+
+
+# TODO(micky774): Update jnp.floor to preserve integral dtype
+def floor(x, /):
+  """Rounds each element x_i of the input array x to the greatest (i.e., closest to +infinity) integer-valued number that is not greater than x_i."""
+  x, = promote_args("floor", x)
+  if isdtype(x.dtype, "integral"):
+    return x
+  return jax.numpy.floor(x)
+
+
+# TODO(micky774): Remove when jnp.hypot deprecation is completed
+# (began 2024-4-14) and default behavior is Array API 2023 compliant
+def hypot(x1, x2, /):
+  """Computes the square root of the sum of squares for each element x1_i of the input array x1 with the respective element x2_i of the input array x2."""
+  x1, x2 = promote_args("hypot", x1, x2)
+
+  if issubdtype(x1.dtype, jax.numpy.complexfloating):
+    raise ValueError(
+      "hypot does not support complex-valued inputs. Please convert to real "
+      "values first, such as by using jnp.real or jnp.imag to take the real "
+      "or imaginary components respectively.")
+  return jax.numpy.hypot(x1, x2)
+
+
+# TODO(micky774): Update jnp.trunc to preserve integral dtype
+def trunc(x, /):
+  """Rounds each element x_i of the input array x to the nearest integer-valued number that is closer to zero than x_i."""
+  x, = promote_args("trunc", x)
+  if isdtype(x.dtype, "integral"):
+    return x
+  return jax.numpy.trunc(x)

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_api/_fft_functions.py

@@ -0,0 +1,25 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import jax.numpy as jnp
+
+# TODO(micky774): Remove after adding device parameter to corresponding jnp.fft
+# functions.
+def fftfreq(n, /, *, d=1.0, device=None):
+  """Returns the discrete Fourier transform sample frequencies."""
+  return jnp.fft.fftfreq(n, d=d).to_device(device)
+
+def rfftfreq(n, /, *, d=1.0, device=None):
+  """Returns the discrete Fourier transform sample frequencies (for rfft and irfft)."""
+  return jnp.fft.rfftfreq(n, d=d).to_device(device)

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_api/_linear_algebra_functions.py

@@ -0,0 +1,28 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import jax
+
+# TODO(micky774): Remove after deprecation is completed (began 2024-5-14)
+def matrix_rank(x, /, *, rtol=None):
+  """
+  Returns the rank (i.e., number of non-zero singular values) of a matrix (or a stack of matrices).
+  """
+  return jax.numpy.linalg.matrix_rank(x, rtol)
+
+def pinv(x, /, *, rtol=None):
+  """
+  Returns the (Moore-Penrose) pseudo-inverse of a matrix (or a stack of matrices) x.
+  """
+  return jax.numpy.linalg.pinv(x, rtol)

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_api/_manipulation_functions.py

@@ -0,0 +1,25 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import jax
+from jax import Array
+
+
+# TODO(micky774): Implement copy
+def reshape(x: Array, /, shape: tuple[int, ...], *, copy: bool | None = None) -> Array:
+  """Reshapes an array without changing its data."""
+  del copy  # unused
+  return jax.numpy.reshape(x, shape)

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_api/_statistical_functions.py

@@ -0,0 +1,25 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import jax
+
+
+def std(x, /, *, axis=None, correction=0.0, keepdims=False):
+  """Calculates the standard deviation of the input array x."""
+  return jax.numpy.std(x, axis=axis, correction=correction, keepdims=keepdims)
+
+
+def var(x, /, *, axis=None, correction=0.0, keepdims=False):
+  """Calculates the variance of the input array x."""
+  return jax.numpy.var(x, axis=axis, correction=correction, keepdims=keepdims)

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_api/_utility_functions.py

@@ -0,0 +1,86 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import jax
+from typing import Tuple
+from jax._src.sharding import Sharding
+from jax._src.lib import xla_client as xc
+from jax._src import dtypes as _dtypes, config
+
+# TODO(micky774): Add to jax.numpy.util when finalizing jax.experimental.array_api
+# deprecation
+class __array_namespace_info__:
+
+  def __init__(self):
+    self._capabilities = {
+      "boolean indexing": True,
+      "data-dependent shapes": False,
+    }
+
+
+  def _build_dtype_dict(self):
+    array_api_types = {
+      "bool", "int8", "int16",
+      "int32", "uint8", "uint16",
+      "uint32", "float32", "complex64"
+    }
+    if config.enable_x64.value:
+      array_api_types |= {"int64", "uint64", "float64", "complex128"}
+    return {category: {t.name: t for t in types if t.name in array_api_types}
+            for category, types in _dtypes._dtype_kinds.items()}
+
+  def default_device(self):
+    # By default JAX arrays are uncommitted (device=None), meaning that
+    # JAX is free to choose the most efficient device placement.
+    return None
+
+  def devices(self):
+    return jax.devices()
+
+  def capabilities(self):
+    return self._capabilities
+
+  def default_dtypes(self, *, device: xc.Device | Sharding | None = None):
+    # Array API supported dtypes are device-independent in JAX
+    del device
+    default_dtypes = {
+      "real floating": "f",
+      "complex floating": "c",
+      "integral": "i",
+      "indexing": "i",
+    }
+    return {
+      dtype_name: _dtypes.canonicalize_dtype(
+        _dtypes._default_types.get(kind)
+      ) for dtype_name, kind in default_dtypes.items()
+    }
+
+  def dtypes(
+      self, *,
+      device: xc.Device | Sharding | None = None,
+      kind: str | Tuple[str, ...] | None = None):
+    # Array API supported dtypes are device-independent in JAX
+    del device
+    data_types = self._build_dtype_dict()
+    if kind is None:
+      out_dict = data_types["numeric"] | data_types["bool"]
+    elif isinstance(kind, tuple):
+      out_dict = {}
+      for _kind in kind:
+        out_dict |= data_types[_kind]
+    else:
+      out_dict = data_types[kind]
+    return out_dict

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_api/_version.py

@@ -0,0 +1,15 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+__array_api_version__ = '2023.12'

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_api/fft.py

@@ -0,0 +1,33 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from jax.numpy.fft import (
+  fft as fft,
+  fftn as fftn,
+  fftshift as fftshift,
+  hfft as hfft,
+  ifft as ifft,
+  ifftn as ifftn,
+  ifftshift as ifftshift,
+  ihfft as ihfft,
+  irfft as irfft,
+  irfftn as irfftn,
+  rfft as rfft,
+  rfftn as rfftn,
+)
+
+from jax.experimental.array_api._fft_functions import (
+  fftfreq as fftfreq,
+  rfftfreq as rfftfreq,
+)

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_api/linalg.py

@@ -0,0 +1,43 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from jax.numpy.linalg import (
+  cholesky as cholesky,
+  cross as cross,
+  det as det,
+  diagonal as diagonal,
+  eigh as eigh,
+  eigvalsh as eigvalsh,
+  inv as inv,
+  matmul as matmul,
+  matrix_norm as matrix_norm,
+  matrix_power as matrix_power,
+  matrix_transpose as matrix_transpose,
+  outer as outer,
+  qr as qr,
+  slogdet as slogdet,
+  solve as solve,
+  svd as svd,
+  svdvals as svdvals,
+  tensordot as tensordot,
+  vecdot as vecdot,
+  vector_norm as vector_norm,
+)
+
+from jax.numpy.linalg import trace as trace
+
+from jax.experimental.array_api._linear_algebra_functions import (
+  matrix_rank as matrix_rank,
+  pinv as pinv,
+)

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_serialization/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2021 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_serialization/serialization.py

@@ -0,0 +1,635 @@
+# Copyright 2021 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Array serialization and deserialization."""
+
+from __future__ import annotations
+
+import abc
+import asyncio
+from collections.abc import Awaitable, Sequence
+from functools import partial
+import itertools
+import logging
+import os
+import re
+import sys
+import threading
+import time
+from typing import Any, Callable, Optional, Union
+
+import jax
+from jax._src import array
+from jax._src import distributed
+from jax._src import sharding
+from jax._src import sharding_impls
+from jax._src.layout import Layout, DeviceLocalLayout as DLL
+from jax._src import typing
+from jax._src import util
+from jax._src.lib import xla_extension as xe
+import jax.numpy as jnp
+import numpy as np
+import tensorstore as ts
+
+
+TS_CONTEXT = ts.Context({'file_io_concurrency': {'limit': 128}})
+_REMOVED_VALUE = 'Value removed'
+_CHECKPOINT_SUCCESS = 'checkpoint_write_success'
+_module_unique_count = itertools.count()
+_DEFAULT_DRIVER = 'file'
+_DISTRIBUTED_SYSTEM_MSG = (
+    'Please initialize the distributed system via '
+    '`jax.distributed.initialize()` at the start of your program.')
+_REMOTE_URL_PREFIXES = ['gs://', 's3://']
+_REMOTE_DRIVER_VALIDATIONS = [
+    {'driver': 'gcs', 'path_regex': None},
+    {'driver': 's3', 'path_regex': None},
+]
+
+class BarrierTimeoutException(Exception):
+  pass
+
+_BARRIER_TIMED_OUT_MSG = (
+    "Suggestions for possible fixes:\n"
+    "* Check the logs to see if one or more processes failed.\n"
+    "* Make sure the training and checkpointing endpoints are close geographically.\n"
+    "* Try increasing the timeout you pass to GlobalAsyncCheckpointManager.")
+
+logger = logging.getLogger(__name__)
+
+
+async def create_async_array_from_callback(
+    global_shape: array.Shape,
+    inp_sharding: jax.sharding.Sharding,
+    data_callback: Callable[[array.Index, jax.Device], Awaitable[jax.Array]],
+):
+  device_to_index_map = inp_sharding.devices_indices_map(global_shape)
+  addressable_da = inp_sharding._addressable_device_assignment
+  future_arrays = [data_callback(device_to_index_map[d], d)
+                   for d in addressable_da]
+  dbs = await asyncio.gather(*future_arrays)
+  return array.make_array_from_single_device_arrays(
+      global_shape, inp_sharding, dbs)
+
+
+def _get_metadata(arr):
+  local_shape = arr.addressable_data(0).shape
+  return {
+      'compressor': {'id': 'zstd'},
+      'shape': arr.shape,
+      'chunks': np.array(np.maximum(1, local_shape)),
+  }
+
+
+def _spec_has_metadata(tree):
+  if not isinstance(tree, dict):
+    return False
+  return 'metadata' in tree or any(
+      _spec_has_metadata(subtree) for _, subtree in tree.items())
+
+def _get_kvstore_for_gcs(ckpt_path: str):
+  m = re.fullmatch('^gs://([^/]*)/(.*)$', ckpt_path, re.DOTALL)
+  if m is None:
+    raise ValueError('The ckpt_path should contain the bucket name and the '
+                      f'file path inside the bucket. Got: {ckpt_path}')
+  gcs_bucket = m.group(1)
+  path_without_bucket = m.group(2)
+  return {'driver': 'gcs', 'bucket': gcs_bucket, 'path': path_without_bucket}
+
+def get_tensorstore_spec(ckpt_path: str, ocdbt: bool = False):
+  # Normalize path to exclude trailing '/'. In GCS path case, we will need to
+  # fix the path prefix to add back the stripped '/'.
+  ckpt_path = os.path.normpath(ckpt_path).replace('gs:/', 'gs://')
+  is_gcs_path = ckpt_path.startswith('gs://')
+  spec = {'driver': 'zarr', 'kvstore': {}}
+  if ocdbt:
+    if not is_gcs_path and not os.path.isabs(ckpt_path):
+      raise ValueError(f'Checkpoint path should be absolute. Got {ckpt_path}')
+    base_path = os.path.dirname(ckpt_path)
+    spec['kvstore'] = {
+        'driver': 'ocdbt',
+        'base': base_path if is_gcs_path else f'{_DEFAULT_DRIVER}://{base_path}',
+        'path': os.path.basename(ckpt_path),
+    }
+  else:
+    if is_gcs_path:
+      spec['kvstore'] = _get_kvstore_for_gcs(ckpt_path)
+    else:
+      spec['kvstore'] = {'driver': _DEFAULT_DRIVER, 'path': ckpt_path}
+
+  return spec
+
+
+def is_remote_storage(tspec: Union[dict[str, Any], str]) -> bool:
+  """Detect if user is using cloud storages.
+
+  This can detect common defines and unable to detect some corner cases such as
+  using gcsfuse.
+  """
+  if isinstance(tspec, str):
+    # KvStoreUrl
+    if re.match(rf'^({"|".join(_REMOTE_URL_PREFIXES)})', tspec):
+      return True
+    else:
+      return False
+
+  for key in ('base', 'kvstore'):
+    if key in tspec:
+      return is_remote_storage(tspec[key])
+
+  if 'driver' in tspec:
+    for rule in _REMOTE_DRIVER_VALIDATIONS:
+      if tspec['driver'] == rule['driver']:
+        if rule['path_regex'] is None:
+          return True
+
+        # check if path matches the regex.
+        if re.match(rule['path_regex'], tspec['path']):
+          return True
+
+  return False
+
+
+# Lifted from T5X.
+class _LimitInFlightBytes:
+  """Limits in-flight bytes when reading/writing checkpoints per process."""
+
+  def __init__(self, num_bytes):
+    self._max_bytes = num_bytes
+    self._available_bytes = num_bytes
+    self._cv = asyncio.Condition(lock=asyncio.Lock())
+
+  async def wait_for_bytes(self, requested_bytes):
+    if requested_bytes >= self._max_bytes:
+      raise ValueError('Requested more bytes than we reserved space for: '
+                       f'{requested_bytes} > {self._max_bytes}')
+    async with self._cv:
+      await self._cv.wait_for(lambda: self._available_bytes > requested_bytes)
+      self._available_bytes -= requested_bytes
+      assert self._available_bytes >= 0
+
+  async def release_bytes(self, requested_bytes):
+    async with self._cv:
+      self._available_bytes += requested_bytes
+      assert self._available_bytes <= self._max_bytes
+      self._cv.notify_all()
+
+
+async def async_serialize(
+    arr_inp,
+    tensorstore_spec,
+    commit_future=None,
+    context=TS_CONTEXT,
+    primary_host: Optional[int] = 0,
+    replica_id: int = 0,
+):
+  """Serialize an array using TensorStore.
+
+  Args:
+    arr_inp: The array to serialize.
+    tensorstore_spec: The tensorstore spec to use.
+    commit_future: A list of futures that will be appended to. The futures can
+      be awaited asynchronously. If None, the futures will be awaited
+      synchronously by this method.
+    context: ts.Context instance.
+    primary_host: Primary host, which indicates the host that will be treated as
+      the "leader". If None, all hosts are treated as the primary. DO NOT USE
+      unless you are sure you know what you are doing.
+    replica_id: Allows overriding the shard replica id that will be saved.
+      DO NOT USE unless you are sure you know what you are doing.
+  """
+  if (isinstance(arr_inp, array.ArrayImpl) and jax.process_count() > 1 and
+      arr_inp.is_fully_addressable):
+    raise ValueError(
+        f'Passing fully addressable arrays to a multiprocess '
+        f'serialization is not allowed, as this may lead to a race condition '
+        f'between processes. Serialization have failed for the array with '
+        f'the path "{tensorstore_spec["kvstore"]["path"]}".')
+
+  # 'metadata' may not be present at the top level (for example, if we are using
+  # a 'cast' driver).
+  if not _spec_has_metadata(tensorstore_spec):
+    tensorstore_spec['metadata'] = _get_metadata(arr_inp)
+
+  # Set dtype if it's not in spec
+  if 'dtype' not in tensorstore_spec:
+    tensorstore_spec['dtype'] = jnp.dtype(arr_inp.dtype).name
+
+  # If primary_host is None, all hosts will checkpoint. This is used
+  # for checkpointing to local filesystem.
+  if primary_host is None or jax.process_index() == primary_host:
+    open_future = ts.open(
+        ts.Spec(tensorstore_spec),
+        create=True,
+        open=True,
+        context=context,
+    )
+    # Asynchronous case.
+    if commit_future is not None:
+      assert isinstance(commit_future, list)
+      commit_future.append(open_future)
+    else:
+      await open_future
+
+  # `ts.open` runs twice for process `primary_host` because for the first time,
+  # we just get the future to be awaited upon in the background thread. The
+  # second one runs with `assume_metadata=True` which does no I/O operation and
+  # returns the tensorstore object.
+  # For every process other than `primary_host`, we open with
+  # `assume_metadata=True`.
+  t = await ts.open(
+      ts.Spec(tensorstore_spec),
+      open=True,
+      assume_metadata=True,
+      context=context,
+  )
+
+  async def _write_array(shard):
+    if shard.replica_id == replica_id:
+      write_future = t[shard.index].write(shard.data)
+      if commit_future is not None:
+        assert isinstance(commit_future, list)
+        commit_future.append(write_future.commit)
+        await write_future.copy
+      else:
+        await write_future.commit
+
+  local_shards = arr_inp.addressable_shards
+  future_write_state = jax.tree_util.tree_map(_write_array, local_shards)
+  return await asyncio.gather(*future_write_state)
+
+
+def run_serialization(arrays, tensorstore_specs):
+  async def _run_serializer():
+    future_writer = jax.tree_util.tree_map(async_serialize, arrays, tensorstore_specs)
+    return await asyncio.gather(*future_writer)
+  asyncio.run(_run_serializer())
+
+
+def estimate_read_memory_footprint(t: ts.TensorStore,
+                                   domain: ts.IndexDomain) -> int:
+  rank = t.rank
+  num_bytes = t.dtype.numpy_dtype.itemsize
+  chunk_template = t.chunk_layout.read_chunk_template
+  if domain is None:
+    domain = t.domain
+  origin = domain.origin
+  shape = domain.shape
+  chunk_origin = chunk_template.origin
+  chunk_shape = chunk_template.shape
+
+  # Some TensorStore drivers are not chunked, e.g. the inline 'array' driver.
+  # For those, instead of returning a near-infinite memory footprint, estimate
+  # the footprint as the entire shape.
+  for i in range(rank):
+    if not chunk_template[i].finite:
+      return domain.size * num_bytes
+
+  # Otherwise, if we have a chunked driver, estimate based on chunk size.
+  for i in range(rank):
+    origin_value = origin[i]
+    chunk_origin_value = chunk_origin[i]
+    chunk_size = chunk_shape[i]
+    lower = origin_value - chunk_origin_value
+    upper = origin_value + shape[i] - chunk_origin_value
+    lower_aligned = lower // chunk_size * chunk_size
+    upper_aligned = -(-upper // chunk_size) * chunk_size
+    num_bytes *= (upper_aligned - lower_aligned)
+
+  return num_bytes
+
+
+async def async_deserialize(
+    user_in_sharding: jax.sharding.Sharding | Layout,
+    tensorstore_spec: ts.Spec | dict[str, Any],
+    global_shape: Sequence[int] | None = None,
+    dtype=None,
+    byte_limiter: _LimitInFlightBytes | None = None,
+    context=TS_CONTEXT,
+    assume_metadata: bool = False,
+):
+  in_sharding = (user_in_sharding.sharding
+                 if isinstance(user_in_sharding, Layout) else user_in_sharding)
+  if not isinstance(in_sharding, jax.sharding.Sharding):
+    raise ValueError(
+        'sharding passed to deserialization should be specified, concrete and'
+        f' an instance of `jax.sharding.Sharding`. Got {in_sharding}')
+  dll = (user_in_sharding.device_local_layout
+         if isinstance(user_in_sharding, Layout) else None)
+  t = await ts.open(
+      tensorstore_spec,
+      open=True,
+      assume_metadata=assume_metadata,
+      context=context,
+  )
+  shape = t.shape if global_shape is None else global_shape
+  new_shard_shape = in_sharding.shard_shape(tuple(shape))
+
+  async def cb(index: array.Index, device: jax.Device):
+    requested_domain = ts.IndexTransform(input_shape=shape)[index].domain
+    restricted_domain = t.domain.intersect(requested_domain)
+    requested_bytes = estimate_read_memory_footprint(t, restricted_domain)
+    # Limit the bytes read for every shard.
+    if byte_limiter is not None:
+      await byte_limiter.wait_for_bytes(requested_bytes)
+    # This maybe needed because the shape the array was saved with is smaller
+    # than the requested shape of the array in which it will be reloaded. So
+    # the extra values will be filled with 0s.
+    out = np.zeros(new_shard_shape, dtype=t.dtype.numpy_dtype)
+    await ts.array(out)[ts.d[:].translate_to[requested_domain.origin]][
+        restricted_domain].write(t[restricted_domain])
+    if dtype is not None:
+      # Cast while reloading on process to avoid 2 copies on device if the
+      # casting is done on device.
+      out = out.astype(dtype)
+    # Convert to jnp array so that layouts are initialized properly for
+    # sub-byte dtypes.
+    # TODO(yashkatariya): This is a band-aid fix. Figure out a better way to
+    # make this work.
+    if out.dtype == jnp.int4:
+      out = jnp.asarray(out)  # type: ignore
+    result = jax.device_put(
+        out, Layout(dll, jax.sharding.SingleDeviceSharding(device)))
+    if byte_limiter is not None:
+      # NB: `out` actually might not be ready for garbage collection by the
+      # time we call release_bytes . Thus peak memory usage still might grow
+      # beyond what byte_limiter limit suggests it should. The simplest option
+      # would be to call  `result.block_until_ready()`` here. However it
+      # also comes with ~15-20% perf penalty as we would be waiting for CPU->GPU
+      # transfer instead of loading data. In the future, if memory pressure
+      # becomes a problem, we can instead instrument  bytelimiter to
+      # keep track of all in-flight tensors and only block_until_ready, if byte
+      # limiter hits the limit to get reduced memory usage, without losing
+      # performance in common use cases.
+      await byte_limiter.release_bytes(requested_bytes)
+    return result
+
+  return await create_async_array_from_callback(tuple(shape), in_sharding, cb)
+
+
+def run_deserialization(shardings: Sequence[sharding.Sharding | Layout],
+                        tensorstore_specs: Sequence[dict[str, Any]],
+                        global_shapes: Sequence[array.Shape] | None = None,
+                        dtypes: Sequence[typing.DTypeLike] | None = None,
+                        concurrent_gb: int = 32):
+  concurrent_bytes = concurrent_gb * 10**9
+
+  async def _run_deserializer():
+    # Object should be created once per process.
+    byte_limiter = _LimitInFlightBytes(concurrent_bytes)
+
+    future_arrays = jax.tree_util.tree_map(
+        partial(async_deserialize, byte_limiter=byte_limiter),
+        shardings, tensorstore_specs,
+        [None] * len(tensorstore_specs) if global_shapes is None else global_shapes,
+        [None] * len(tensorstore_specs) if dtypes is None else dtypes)
+    return await asyncio.gather(*future_arrays)
+  return asyncio.run(_run_deserializer())
+
+
+def _get_key(key: int):
+  return f'tensorstore_checkpoint_{key}'
+
+
+class GlobalAsyncCheckpointManagerBase(util.StrictABC):
+  """Interface for checkpointing GDAs asynchronously.
+
+  This class manages the state of an ongoing asynchronous checkpoint.
+
+  For example, say a checkpoint happens on every step. If you checkpoint on
+  step 1 and after some computation the model is on checkpoint 2. But step 1's
+  checkpoint hasn't finished committing to the storage layer yet. So until that
+  is finished, checkpoint for step 2 will need to be blocked. Maintaining a
+  class allows to maintain that state.
+
+  Example:
+
+  Below is a simplified training loop:
+
+  ```
+  # Call this at the start of your program.
+  jax.distributed.initialize()
+
+  manager = GlobalAsyncCheckpointManager()
+
+  # Restore checkpoint if available or initialize the train_state from
+  # init_fn().
+  train_state = manager.deserialize(...)
+
+  while ...:
+    if step % num_steps_between_checkpoints == 0:
+      manager.serialize(train_state, temp_checkpoint_dir=...,
+                        final_checkpoint_dir=...)
+      train_state = train_step(train_state, input)
+      # This is a non-blocking call.
+      manager.check_for_errors()
+
+  manager.serialize(train_state, temp_checkpoint_dir=...,
+                    final_checkpoint_dir=...)
+  # Wait before the end of the program for the checkpoint to finish. This is a
+  # blocking call.
+  manager.wait_until_finished()
+  ```
+  """
+
+  @abc.abstractmethod
+  def check_for_errors(self):
+    """Checks if any errors have been raised in the child thread.
+
+    This is a non-blocking call that can be called in the main thread.
+    """
+
+  @abc.abstractmethod
+  def wait_until_finished(self):
+    """Blocks until serialization has finished."""
+
+  @abc.abstractmethod
+  def serialize(self, arrays, tensorstore_specs, *,
+                on_commit_callback: Callable[[], None]):
+    """Serializes GDAs to TensorStore."""
+
+  @abc.abstractmethod
+  def deserialize(self, shardings: Sequence[sharding.Sharding],
+                  tensorstore_specs: Sequence[dict[str, Any]],
+                  global_shapes: Sequence[array.Shape] | None = None,
+                  dtypes: Sequence[typing.DTypeLike] | None = None):
+    """Deserializes GDAs from TensorStore."""
+
+
+class AsyncManager:
+
+  def __init__(self, timeout_secs=300):
+    self._timeout_secs = timeout_secs
+    self._timeout_in_ms = self._timeout_secs * 1000
+
+    self._commit_futures = None
+    self._thread = None
+    self._exception = None
+
+    if jax.process_count() > 1 and distributed.global_state.client is None:
+      raise ValueError(_DISTRIBUTED_SYSTEM_MSG)
+    if jax.process_count() > 1:
+      self._client = distributed.global_state.client
+    self._count = None
+
+  def __del__(self):
+    if self._thread is not None and self._thread.is_alive():
+      logger.warning('Please add `.wait_until_finished()` in the main thread '
+                      'before your program finishes because there is a '
+                      'possibility of losing errors raised if the '
+                      'this class is deleted before writing is completed.')
+
+  def _thread_func(self):
+    try:
+      current_process = jax.process_index()
+      process_count = jax.process_count()
+      logger.info('Starting commit to storage layer by process: %s',
+                   current_process)
+      thread_start_time = time.time()
+      for future in self._commit_futures:
+        future.result()
+      logger.info('Finished committing to storage layer by process: %s',
+                   current_process)
+
+      if process_count > 1:
+        # All processes will wait at the barrier. When all processes are at the
+        # barrier, the barrier will be satisfied. If not, then it will timeout.
+        key_for_barrier = _get_key(self._count)
+        logger.info('Key used for barrier is %s for process %s',
+                    key_for_barrier, current_process)
+        self._client.wait_at_barrier(key_for_barrier, self._timeout_in_ms)
+        logger.info('Finished waiting at barrier for process %s',
+                    current_process)
+
+      if current_process == 0:
+        self._on_commit_callback()
+        logger.info('on_commit_callback successfully ran!')
+        if process_count > 1:
+          self._client.key_value_set(key_for_barrier, _CHECKPOINT_SUCCESS)
+          logger.info('Process 0 successfully set key %s in the kv store',
+                      key_for_barrier)
+
+      jax.monitoring.record_event_duration_secs(
+          '/jax/checkpoint/write/async/thread_duration_sec',
+          time.time() - thread_start_time)
+
+    except Exception as e:
+      self._exception = e
+
+  def _start_async_commit(self, on_commit_callback):
+    self._count = next(_module_unique_count)
+
+    self._on_commit_callback = on_commit_callback
+    self._thread = threading.Thread(target=self._thread_func)
+    self._thread.start()
+
+  def check_for_errors(self):
+    if self._exception is not None:
+      # Clears self._exception so it is only raised once.
+      exception = self._exception
+      self._exception = None
+      if (isinstance(exception, xe.XlaRuntimeError) and
+          'DEADLINE_EXCEEDED: Barrier timed out' in str(exception)):
+        raise BarrierTimeoutException(
+            '\n'.join([str(exception), _BARRIER_TIMED_OUT_MSG]))
+      raise exception  # pylint: disable=raising-bad-type
+
+  def wait_until_finished(self):
+    if self._thread is not None:
+      self._thread.join()
+      self._thread = None
+      logger.info('Thread joined successfully')
+
+    self.check_for_errors()
+    logger.info('Error check finished successfully')
+
+    if jax.process_count() > 1 and self._count is not None:
+      # Block until process 0 writes success value to the key value store.
+      # If it fails to write it, then `blocking_key_value_get` will time out.
+      get_key = _get_key(self._count)
+      self._client.blocking_key_value_get(get_key, self._timeout_in_ms)
+      logger.info('blocking_key_value_get on key %s was successfully '
+                  'completed.', get_key)
+
+  def _add_futures(self, futures: Sequence[asyncio.Future]):
+    self._commit_futures = futures
+
+
+class GlobalAsyncCheckpointManager(AsyncManager, GlobalAsyncCheckpointManagerBase):
+  """Responsible for serializing GDAs via TensorStore."""
+
+  def serialize(self, arrays, tensorstore_specs, *, on_commit_callback):
+    """Serializes Arrays or Arrays via TensorStore asynchronously.
+
+    TensorStore writes to a storage layer in 2 steps:
+    *  Reading/copying from the source after which the source can be modified.
+         * Returns a copy future.
+    *  Writing/committing to the storage layer.
+         * Returns a commit future.
+
+    In asynchronous mode, the serialization waits for the commit future to
+    finish in a separate thread allowing other computation to proceed.
+
+    Args:
+      arrays: Arrays or Arrays that should be serialized.
+      tensorstore_specs: TensorStore specs that are used to serialize GDAs or
+        Arrays.
+      on_commit_callback: This callback will be executed after all processes
+        have finished writing their checkpoints to disk. Filesystems where
+        atomic rename operations are supported, you can rename from the
+        temporary directory to the final directory. On GCS, you write to the
+        final directory directly and in `on_commit_callback` you write a
+        success file indicating that the serialization was successful because
+        GCS does not support atomic rename operations.
+    """
+    logger.info('Waiting for previous serialization to finish.')
+    self.wait_until_finished()
+
+    commit_futures = [[] for _ in range(len(tensorstore_specs))]
+
+    async def _run_serializer():
+      future_writer = jax.tree_util.tree_map(
+          async_serialize, arrays, tensorstore_specs, commit_futures)
+      return await asyncio.gather(*future_writer)
+
+    asyncio.run(_run_serializer())
+
+    self._add_futures(jax.tree_util.tree_flatten(commit_futures)[0])
+
+    # Used in wait_until_finished to check on process != 0, if the checkpoint
+    # has finished writing.
+    self._start_async_commit(on_commit_callback)
+
+  def serialize_with_paths(self, arrays: Sequence[jax.Array],
+                           paths: Sequence[str], *, on_commit_callback):
+    tspecs = jax.tree.map(get_tensorstore_spec, paths)
+    self.serialize(arrays, tspecs, on_commit_callback=on_commit_callback)
+
+  def deserialize(self, shardings: Sequence[sharding.Sharding | Layout],
+                  tensorstore_specs: Sequence[dict[str, Any]],
+                  global_shapes: Sequence[array.Shape] | None = None,
+                  dtypes: Sequence[typing.DTypeLike] | None = None,
+                  concurrent_gb: int = 32):
+    self.wait_until_finished()
+    return run_deserialization(shardings, tensorstore_specs,
+                               global_shapes, dtypes, concurrent_gb)
+
+  def deserialize_with_paths(
+      self, shardings: Sequence[sharding.Sharding],
+      paths: Sequence[str],
+      global_shapes: Sequence[array.Shape] | None = None,
+      dtypes: Sequence[typing.DTypeLike] | None = None,
+      concurrent_gb: int = 32):
+    tspecs = jax.tree.map(get_tensorstore_spec, paths)
+    return self.deserialize(shardings, tspecs, global_shapes, dtypes,
+                            concurrent_gb)

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/array_serialization/serialization_test.py

@@ -0,0 +1,493 @@
+# Copyright 2021 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for serialization and deserialization of GDA."""
+
+import asyncio
+import contextlib
+import math
+from functools import partial
+import re
+import os
+import pathlib
+import tracemalloc as tm
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+import jax.numpy as jnp
+from jax._src import test_util as jtu
+from jax._src import array
+from jax._src import xla_bridge as xb
+from jax.sharding import NamedSharding, GSPMDSharding
+from jax.sharding import PartitionSpec as P
+from jax.experimental.array_serialization import serialization
+from jax.experimental.layout import Layout, DeviceLocalLayout as DLL
+import numpy as np
+import tensorstore as ts
+
+jax.config.parse_flags_with_absl()
+_exit_stack = contextlib.ExitStack()
+
+def setUpModule():
+  _exit_stack.enter_context(jtu.set_host_platform_device_count(8))
+
+def tearDownModule():
+  _exit_stack.close()
+
+
+pattern = re.compile(r"\{(.*?):")
+
+def extract_minor_to_major(l):
+  match = re.search(pattern, str(l))
+  return tuple(int(i) for i in match.groups()[0].split(','))
+
+
+class CheckpointTest(jtu.JaxTestCase):
+
+  def _on_commit_callback(self, temp_ckpt_dir, final_ckpt_dir):
+    os.rename(temp_ckpt_dir, final_ckpt_dir)
+
+  @jtu.skip_on_devices('cpu')
+  def test_memory_consumption(self):
+    global_mesh = jtu.create_global_mesh((2, 4), ('x', 'y'))
+    inp_shape = (2_048, 4_096)
+    pspec = P('x', 'y')
+    num = math.prod(inp_shape)
+    sharding = NamedSharding(global_mesh, pspec)
+    src = jnp.arange(num, dtype=np.int32).reshape(inp_shape)  # 8e9
+    inp = array.make_array_from_callback(
+        inp_shape, sharding,
+        lambda idx: src[idx])
+    ckpt_dir = pathlib.Path(self.create_tempdir('memprof').full_path)
+    tspec = serialization.get_tensorstore_spec(str(ckpt_dir))
+
+    manager = serialization.GlobalAsyncCheckpointManager()
+    manager.serialize(
+        [inp], [tspec],
+        on_commit_callback=partial(self._on_commit_callback, ckpt_dir, ckpt_dir))
+    manager.wait_until_finished()
+
+    async def deserialize_with_byte_limit():
+      r = await serialization.async_deserialize(
+        sharding, tspec, inp_shape,
+        byte_limiter=serialization._LimitInFlightBytes(4_200_000))
+      r.block_until_ready()
+
+    tm.start()
+    asyncio.run(deserialize_with_byte_limit())
+    unused_current, peak = tm.get_traced_memory()
+    # NB: some padding + tensorstore overhead. It should always be
+    # less than array size (2048 * 4096 * 4 = 32M)
+    self.assertLess(peak, 10_000_000)
+    deserialize_wo_limit = serialization.async_deserialize(
+        sharding, tspec, inp_shape)
+    tm.clear_traces()
+    # NB: call block_until_ready() is important here and above
+    # because otherwise this leads to racing condition and segfault with
+    # tensorstore attempting to dealloc using tracemalloc which is already
+    # destroyed.
+    asyncio.run(deserialize_wo_limit).block_until_ready()
+
+    unused_current, peak = tm.get_traced_memory()
+    # We load entire array in memory here.
+    self.assertGreater(peak, 30_000_000)
+    tm.stop()
+
+  def test_checkpointing_with_path_variant(self):
+    global_mesh = jtu.create_global_mesh((4, 2), ('x', 'y'))
+    inp_shape = (8, 2)
+    pspec = P('x', 'y')
+    num = math.prod(inp_shape)
+
+    # First Array
+    global_input_data1 = np.arange(num, dtype=np.int32).reshape(inp_shape)
+    a1 = array.make_array_from_callback(
+        inp_shape, NamedSharding(global_mesh, pspec),
+        lambda idx: global_input_data1[idx])
+    ckpt_dir = pathlib.Path(self.create_tempdir('ckpt_variant').full_path)
+    ckpt_path1 = pathlib.Path(
+        self.create_tempdir(f'{ckpt_dir}/first').full_path)
+
+    ckpt_paths = [str(ckpt_path1)]
+    manager = serialization.GlobalAsyncCheckpointManager()
+    manager.serialize_with_paths(
+        [a1], ckpt_paths,
+        on_commit_callback=partial(self._on_commit_callback, ckpt_dir, ckpt_dir))
+    manager.wait_until_finished()
+
+    m1, = manager.deserialize_with_paths(
+        [NamedSharding(global_mesh, pspec)], ckpt_paths)
+    self.assertIsInstance(m1, array.ArrayImpl)
+    self.assertArraysEqual(np.asarray(m1.addressable_shards[0].data),
+                           np.array([[0], [2]], dtype=np.int32))
+    self.assertArraysEqual(np.asarray(m1.addressable_shards[1].data),
+                           np.array([[1], [3]], dtype=np.int32))
+    self.assertEqual(m1.addressable_shards[0].data.shape, (2, 1))
+    self.assertEqual(m1.dtype, np.int32)
+
+  def test_checkpointing_jax_array(self):
+    global_mesh = jtu.create_global_mesh((4, 2), ('x', 'y'))
+    inp_shape = (8, 2)
+    pspec = P('x', 'y')
+    num = math.prod(inp_shape)
+
+    # First Array
+    global_input_data1 = np.arange(num, dtype=np.int32).reshape(inp_shape)
+    a1 = array.make_array_from_callback(
+        inp_shape, NamedSharding(global_mesh, pspec),
+        lambda idx: global_input_data1[idx])
+    ckpt_dir = pathlib.Path(self.create_tempdir('ckpt').full_path)
+    ckpt_path1 = pathlib.Path(self.create_tempdir(f'{ckpt_dir}/first').full_path)
+
+    # Second Array
+    global_input_data2 = np.arange(
+        num, num + num, dtype=np.int32).reshape(inp_shape)
+    a2 = array.make_array_from_callback(
+        inp_shape, NamedSharding(global_mesh, pspec),
+        lambda idx: global_input_data2[idx])
+    ckpt_path2 = pathlib.Path(self.create_tempdir(f'{ckpt_dir}/second').full_path)
+
+    # Third Array
+    def cb3(_):
+      return np.array([], dtype=np.float32)
+    global_mesh1d = jtu.create_global_mesh((8,), ('x',))
+    a3 = array.make_array_from_callback(
+        (0,), NamedSharding(global_mesh1d, P(None)), cb3)
+    ckpt_path3 = pathlib.Path(self.create_tempdir(f'{ckpt_dir}/third').full_path)
+
+    ckpt_paths = [str(ckpt_path1), str(ckpt_path2), str(ckpt_path3)]
+    tspecs = jax.tree_util.tree_map(serialization.get_tensorstore_spec, ckpt_paths)
+
+    manager = serialization.GlobalAsyncCheckpointManager()
+    manager.serialize(
+        [a1, a2, a3], tspecs,
+        on_commit_callback=partial(self._on_commit_callback, ckpt_dir, ckpt_dir))
+    manager.wait_until_finished()
+
+    m1, m2, m3 = serialization.run_deserialization(
+        [NamedSharding(global_mesh, pspec),
+         NamedSharding(global_mesh, P('x')),
+         NamedSharding(global_mesh1d, P(None))],
+        tspecs)
+
+    self.assertIsInstance(m1, array.ArrayImpl)
+    self.assertArraysEqual(np.asarray(m1.addressable_shards[0].data),
+                           np.array([[0], [2]], dtype=np.int32))
+    self.assertArraysEqual(np.asarray(m1.addressable_shards[1].data),
+                           np.array([[1], [3]], dtype=np.int32))
+    self.assertEqual(m1.addressable_shards[0].data.shape, (2, 1))
+    self.assertEqual(m1.dtype, np.int32)
+
+    self.assertIsInstance(m2, array.ArrayImpl)
+    self.assertArraysEqual(np.asarray(m2.addressable_shards[0].data),
+                           np.array([[16, 17], [18, 19]], dtype=np.int32))
+    self.assertArraysEqual(np.asarray(m2.addressable_shards[1].data),
+                           np.array([[16, 17], [18, 19]], dtype=np.int32))
+    self.assertEqual(m2.addressable_shards[0].data.shape, (2, 2))
+    self.assertEqual(m2.dtype, np.int32)
+
+    self.assertIsInstance(m3, array.ArrayImpl)
+    for i, s in enumerate(m3.addressable_shards):
+      self.assertEqual(s.index, (slice(None),))
+      self.assertEqual(s.replica_id, i)
+      self.assertArraysEqual(np.asarray(s.data), np.array([], dtype=np.float32))
+    self.assertEqual(m3.dtype, np.float32)
+
+  @parameterized.product(input_dtype=[np.int32, jnp.bfloat16])
+  def test_checkpointing_with_bigger_shape_jax_array(self, input_dtype):
+    global_mesh = jtu.create_global_mesh((2, 2), ('x', 'y'))
+    global_input_shape = (8, 2)
+    num = math.prod(global_input_shape)
+
+    global_input_data1 = np.arange(num, dtype=input_dtype).reshape(
+        global_input_shape
+    )
+    def cb1(index):
+      return global_input_data1[index]
+    arr = array.make_array_from_callback(
+        global_input_shape, NamedSharding(global_mesh, P('x', 'y')), cb1)
+    ckpt_dir = pathlib.Path(self.create_tempdir('first').full_path)
+
+    ckpt_paths = [str(ckpt_dir)]
+    tspecs = jax.tree_util.tree_map(serialization.get_tensorstore_spec, ckpt_paths)
+
+    manager = serialization.GlobalAsyncCheckpointManager()
+    manager.serialize(
+        [arr], tspecs,
+        on_commit_callback=partial(self._on_commit_callback, ckpt_dir, ckpt_dir))
+    manager.wait_until_finished()
+
+    ds = NamedSharding(jtu.create_global_mesh((4, 2), ('x', 'y')), P('x', 'y'))
+
+    m1, = serialization.run_deserialization([ds], tspecs, [(12, 2)],
+                                            [np.float32])
+
+    expected_data = {
+        0: np.array([[0], [2], [4]], dtype=np.float32),
+        1: np.array([[1], [3], [5]], dtype=np.float32),
+        2: np.array([[6], [8], [10]], dtype=np.float32),
+        3: np.array([[7], [9], [11]], dtype=np.float32),
+        4: np.array([[12], [14], [0]], dtype=np.float32),
+        5: np.array([[13], [15], [0]], dtype=np.float32),
+        6: np.array([[0], [0], [0]], dtype=np.float32),
+        7: np.array([[0], [0], [0]], dtype=np.float32),
+    }
+
+    for l in m1.addressable_shards:
+      self.assertArraysEqual(np.asarray(l.data), expected_data[l.device.id])
+
+    new_ds = GSPMDSharding.get_replicated(list(global_mesh.devices.flat))
+    m2, = serialization.run_deserialization([new_ds], tspecs, [(8, 2)], [np.float32])
+    for l in m2.addressable_shards:
+      self.assertArraysEqual(l.data, global_input_data1.astype('float32'))
+
+  @parameterized.product(input_dtype=[jnp.int4, jnp.int8])
+  def test_checkpointing_with_int4(self, input_dtype):
+    global_mesh = jtu.create_global_mesh((2, 2), ('x', 'y'))
+    global_input_shape = (8, 2)
+    num = math.prod(global_input_shape)
+
+    global_input_data = np.arange(num, dtype=input_dtype).reshape(
+        global_input_shape
+    )
+    def cb(index):
+      return global_input_data[index]
+    arr = array.make_array_from_callback(
+        global_input_shape, NamedSharding(global_mesh, P('x', 'y')), cb)
+    ckpt_dir = pathlib.Path(self.create_tempdir('first').full_path)
+
+    ckpt_paths = [str(ckpt_dir)]
+    tspecs = jax.tree_util.tree_map(serialization.get_tensorstore_spec, ckpt_paths)
+
+    manager = serialization.GlobalAsyncCheckpointManager()
+    manager.serialize(
+        [arr], tspecs,
+        on_commit_callback=partial(self._on_commit_callback, ckpt_dir, ckpt_dir))
+    manager.wait_until_finished()
+
+    ds = NamedSharding(jtu.create_global_mesh((4, 2), ('x', 'y')), P('x', 'y'))
+
+    target_dtype = jnp.dtype('int4')
+    m1, = serialization.run_deserialization([ds], tspecs, [(12, 2)],
+                                            [target_dtype])
+
+    # values bigger than 7 are converted properly.
+    expected_data = {
+        0: jnp.array([[0], [2], [4]], dtype=target_dtype),
+        1: jnp.array([[1], [3], [5]], dtype=target_dtype),
+        2: jnp.array([[6], [8], [10]], dtype=target_dtype),
+        3: jnp.array([[7], [9], [11]], dtype=target_dtype),
+        4: jnp.array([[12], [14], [0]], dtype=target_dtype),
+        5: jnp.array([[13], [15], [0]], dtype=target_dtype),
+        6: jnp.array([[0], [0], [0]], dtype=target_dtype),
+        7: jnp.array([[0], [0], [0]], dtype=target_dtype),
+    }
+
+    for l in m1.addressable_shards:
+      self.assertArraysEqual(np.asarray(l.data), expected_data[l.device.id])
+
+    new_ds = GSPMDSharding.get_replicated(list(global_mesh.devices.flat))
+    m2, = serialization.run_deserialization([new_ds], tspecs, [(8, 2)], [target_dtype])
+    for l in m2.addressable_shards:
+      self.assertArraysEqual(l.data, global_input_data.astype(target_dtype))
+
+  def test_checkpointing_scalar_jax_array(self):
+    global_mesh = jtu.create_global_mesh((2,), ('x'))
+    global_input_shape = ()
+    data = np.array(4)
+    s = NamedSharding(global_mesh, P(None))
+    array1 = array.make_array_from_callback(
+        global_input_shape, s, lambda idx: data[idx])
+    ckpt_dir = pathlib.Path(self.create_tempdir('first').full_path)
+
+    ckpt_paths = [str(ckpt_dir)]
+    tspecs = jax.tree_util.tree_map(serialization.get_tensorstore_spec, ckpt_paths)
+
+    manager = serialization.GlobalAsyncCheckpointManager()
+    manager.serialize(
+        [array1], tspecs,
+        on_commit_callback=partial(self._on_commit_callback, ckpt_dir, ckpt_dir))
+    manager.wait_until_finished()
+
+    ds = NamedSharding(jtu.create_global_mesh((2,), ('x')), P(None))
+
+    m1, = serialization.run_deserialization(
+        [ds],
+        tspecs,
+        [()],
+        [np.float32]
+    )
+
+    for l in m1.addressable_shards:
+      self.assertArraysEqual(np.asarray(l.data), data.astype(np.float32))
+
+  def test_deserialize_tensorstore_array_jax_array(self):
+    global_mesh = jtu.create_global_mesh((2,), ('x'))
+    data = np.arange(1024)
+    tspec = ts.array(data).spec()
+    m1, = serialization.run_deserialization(
+        [NamedSharding(global_mesh, P(None))],
+        [tspec]
+    )
+    for l in m1.addressable_shards:
+      self.assertArraysEqual(np.asarray(l.data), data)
+
+  def test_spec_has_metadata(self):
+    spec = {
+        'a': {
+            'b': 1,
+            'c': 2,
+        },
+        'd': 3,
+        'e': {
+            'a': 2,
+            'metadata': 3
+        },
+        'f': 4
+    }
+    self.assertTrue(serialization._spec_has_metadata(spec))
+    self.assertTrue(
+        serialization._spec_has_metadata({
+            'driver': 'zarr',
+            'kvstore': 'gfile',
+            'metadata': {
+                'chunks': 4,
+                'shape': (32, 64)
+            },
+            'one_more': 'thing'
+        }))
+
+  def test_spec_has_no_metadata(self):
+    spec = {
+        'a': {
+            'b': 1,
+            'c': 2,
+        },
+        'd': 3,
+        'e': {
+            'a': 2,
+        },
+        'f': 4
+    }
+    self.assertFalse(serialization._spec_has_metadata(spec))
+
+  def test_empty_spec_has_no_metadata(self):
+    spec = {}
+    self.assertFalse(serialization._spec_has_metadata(spec))
+
+  @parameterized.named_parameters(
+      ('gcs', 'gs://my/ckpt/dir/path'),
+      ('file', '/my/ckpt/dir/path')
+  )
+  def test_get_tensorstore_spec_ocdbt(self, path):
+    spec = serialization.get_tensorstore_spec(path, ocdbt=True)
+    is_gcs_path = path.startswith('gs://')
+    if is_gcs_path:
+      self.assertEqual(spec['kvstore']['base'], os.path.dirname(path))
+    else:
+      self.assertEqual(spec['kvstore']['base'],
+                       f'{serialization._DEFAULT_DRIVER}://{os.path.dirname(path)}')
+    self.assertEqual(spec['kvstore']['path'], 'path')
+
+  def test_get_tensorstore_spec_not_absolute_path(self):
+    path = 'my/ckpt/path'
+    with self.assertRaisesRegex(ValueError,
+                                "Checkpoint path should be absolute"):
+      serialization.get_tensorstore_spec(path, ocdbt=True)
+
+  def test_maybe_cloud_storage(self):
+    gs_path = 'gs://some-buck/path'
+    gs_spec = serialization.get_tensorstore_spec(gs_path, ocdbt=True)
+    self.assertTrue(serialization.is_remote_storage(gs_spec))
+
+    local_path = '/tmp/checkpoint'
+    local_spec = serialization.get_tensorstore_spec(local_path, ocdbt=True)
+    self.assertFalse(serialization.is_remote_storage(local_spec))
+
+    nested_tspec = {
+        'driver': 'cast',
+        'dtype': 'int32',
+        'base': {
+            'driver': 'zarr',
+            'kvstore': {'driver': 'ocdbt', 'base': 's3://some-bucket/path'},
+        },
+    }
+    self.assertTrue(serialization.is_remote_storage(nested_tspec))
+
+  def test_load_with_layout(self):
+    if not jtu.test_device_matches(['tpu']):
+      self.skipTest('Layouts are only supported on TPUs')
+
+    mesh = jtu.create_global_mesh((4, 2), ('x', 'y'))
+    np_inp = np.arange(32).reshape(8, 4)
+    s = NamedSharding(mesh, P('x', 'y'))
+    arr = jax.device_put(np_inp, s)
+
+    out_layout = jax.jit(lambda x: x.T, out_shardings=Layout(DLL.AUTO)).lower(
+        arr).compile().output_layouts()
+    self.assertEqual(extract_minor_to_major(arr.layout),
+                     extract_minor_to_major(out_layout)[::-1])
+
+    ckpt_dir = pathlib.Path(self.create_tempdir('ckpt').full_path)
+    ckpt_path = pathlib.Path(self.create_tempdir(f'{ckpt_dir}/first').full_path)
+    tspecs = jax.tree_util.tree_map(serialization.get_tensorstore_spec, [ckpt_path])
+
+    manager = serialization.GlobalAsyncCheckpointManager()
+    manager.serialize(
+        [arr], tspecs,
+        on_commit_callback=partial(self._on_commit_callback, ckpt_dir, ckpt_dir))
+    manager.wait_until_finished()
+
+    out, = serialization.run_deserialization([out_layout], tspecs)
+
+    self.assertEqual(out.layout, out_layout)
+    self.assertIsInstance(out, array.ArrayImpl)
+    self.assertArraysEqual(out, np_inp)
+    for s in out.addressable_shards:
+      self.assertArraysEqual(s.data, np_inp[s.index])
+
+  def test_deserialization_with_int4(self):
+    dtype = jnp.int4
+    shape = (8, 2)
+    arr = jnp.arange(np.prod(shape)).reshape(shape).astype(dtype)
+
+    ckpt_dir = pathlib.Path(self.create_tempdir('test_ckpt').full_path)
+
+    # Run serialization.
+    sharding = jax.sharding.GSPMDSharding.get_replicated(jax.devices())
+    tspecs = jax.tree_util.tree_map(
+        serialization.get_tensorstore_spec, [ckpt_dir]
+    )
+    manager = serialization.GlobalAsyncCheckpointManager()
+    manager.serialize(
+        [arr],
+        tspecs,
+        on_commit_callback=lambda: None,
+    )
+    manager.wait_until_finished()
+
+    # Run deserialization.
+    deserialized_arr, = serialization.run_deserialization(
+        shardings=[sharding],
+        tensorstore_specs=tspecs,
+        global_shapes=[shape],
+        dtypes=[dtype],
+    )
+
+    out = deserialized_arr.astype(jnp.int8)  # doesn't crash
+    self.assertEqual(out.dtype, jnp.int8)
+    self.assertArraysEqual(out + out, out * 2)
+
+if __name__ == '__main__':
+  absltest.main(testLoader=jtu.JaxTestLoader())

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/compilation_cache/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2021 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/compilation_cache/compilation_cache.py

@@ -0,0 +1,20 @@
+# Copyright 2021 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from jax._src.compilation_cache import (
+  is_initialized as is_initialized,  # deprecated
+  initialize_cache as initialize_cache,  # deprecated; use set_cache_dir instead
+  set_cache_dir as set_cache_dir,
+  reset_cache as reset_cache,
+)

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/export/__init__.py

@@ -0,0 +1,36 @@
+# Copyright 2023 The JAX Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from jax._src.export._export import (
+    minimum_supported_serialization_version,
+    maximum_supported_serialization_version,
+    Exported,
+    export,
+    call_exported,  # TODO: deprecate
+    call,
+    DisabledSafetyCheck,
+    default_lowering_platform,
+)
+from jax._src.export.shape_poly import (
+    is_symbolic_dim,
+    symbolic_shape,
+    symbolic_args_specs,
+    SymbolicScope,
+)
+from jax._src.export.serialization import (
+    serialize,
+    deserialize,
+)
+from jax._src.export import shape_poly_decision

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/__init__.py

@@ -0,0 +1,23 @@
+# Copyright 2020 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from jax.experimental.jax2tf.jax2tf import (
+  convert as convert,
+  eval_polymorphic_shape as eval_polymorphic_shape,
+  dtype_of_val as dtype_of_val,
+  split_to_logical_devices as split_to_logical_devices,
+  DisabledSafetyCheck as DisabledSafetyCheck,
+  PolyShape as PolyShape  # TODO: deprecate
+)
+from jax.experimental.jax2tf.call_tf import call_tf as call_tf

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/call_tf.py

@@ -0,0 +1,682 @@
+# Copyright 2021 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Allows JAX to call TensorFlow functions with support for autodiff.
+
+**Experimental: please give feedback, and expect changes.**
+
+This module introduces the function :func:`call_tf` that allows JAX to call
+TensorFlow functions.
+
+For examples and details, see
+https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#calling-tensorflow-functions-from-jax.
+
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+import dataclasses
+import functools
+from typing import Any, Callable, Optional
+
+from absl import logging
+import jax
+from jax import dlpack
+from jax import dtypes
+from jax import numpy as jnp
+from jax import tree_util
+from jax._src import ad_util
+from jax._src import core
+from jax._src import effects
+from jax._src import util
+from jax._src.lib import xla_client
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import func as func_dialect
+from jax._src.lib.mlir.dialects import hlo
+from jax.experimental.jax2tf import jax2tf as jax2tf_internal
+from jax.interpreters import mlir
+import numpy as np
+import tensorflow as tf
+
+
+map = util.safe_map
+zip = util.safe_zip
+
+TfConcreteFunction = Any
+TfVal = jax2tf_internal.TfVal
+
+# The platforms for which to use DLPack to avoid copying (only works on GPU
+# and CPU at the moment, and only for Array). For CPU we don't need
+# DLPack, if we are careful.
+_DLPACK_PLATFORMS = ("gpu",)
+
+class UnspecifiedOutputShapeDtype:
+  pass
+
+def call_tf(
+    callable_tf: Callable,
+    has_side_effects=True,
+    ordered=False,
+    output_shape_dtype=UnspecifiedOutputShapeDtype(),
+    call_tf_graph=False,
+) -> Callable:
+  """Calls a TensorFlow function from JAX, with support for reverse autodiff.
+
+  The ``callable_tf`` will be called with TensorFlow-compatible arguments (
+  numpy.ndarray, ``tf.Tensor`` or ``tf.Variable``) or pytrees thereof. The
+  function must return the same type of results.
+
+  If ``call_tf`` appears in a JAX staging context (:func:`jax.jit`,
+  or :func:`jax.pmap`, or :func:`jax.xmap`, or a control-flow primitive) then
+  ``callable_tf`` will be compiled with ``tf.function(callable_tf,
+  jit_compile=True)``
+  and the resulting XLA computation will be embedded in JAX's XLA computation.
+
+  If ``call_tf`` appears outside a JAX staging context, it will be called inline
+  using TensorFlow eager mode.
+
+  The ``call_tf`` supports JAX's reverse-mode autodiff, in which case the
+  ``callable_tf`` will be differentiated using ``tf.GradientTape``. This means
+  that the gradient will be TensorFlow-accurate, e.g., will respect the
+  custom gradients that may be defined for the code in ``callable_tf``.
+
+  For an example and more details see the
+  `README
+  <https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#calling-tensorflow-functions-from-jax>`_.
+
+  Args:
+    callable_tf: a TensorFlow Callable that can take a pytree of TensorFlow
+      arguments.
+    has_side_effects: if True then it ensures that instances of this primitive
+      are not removed or replicated by JAX optimizations such as dead-code
+      elimination.
+    ordered: If true, calls are modeled as having ordered effects.
+    output_shape_dtype: An optional declaration of the expected shape and dtype
+      of the result of the called TensorFlow function. If given it will be used
+      during JAX tracing to form the abstract values of the results of the
+      `call_tf`. If not given then we form a `tf.Graph` for the called
+      TensorFlow function and we use the TensorFlow-inferred shapes and types.
+      Must be a pytree matching the structure of the nested structure returned
+      from the TensorFlow function, containing objects with `.shape` and
+      `.dtype` attributes, e.g., `jax.ShapeDtypeStruct` or `jax.Array`.
+    call_tf_graph: EXPERIMENTAL, DO NOT USE. We may change the name in the
+      future.
+
+  Returns: a JAX callable that can be invoked with JAX pytree arguments, in
+    op-by-op mode or in a staged context. This callable can be used with JAX's
+    reverse-mode autodiff (:func:`jax.grad`).
+  """
+  @jax.custom_vjp
+  def make_call(*args_jax):
+    """We wrap it all in `make_call` so that we can attach custom VJP."""
+
+    args_flat_jax, args_treedef = tree_util.tree_flatten(args_jax)
+    # Canonicalize the arguments; e.g., makes them x32 if JAX is in 32-bit mode
+    def canonical_arg(v):
+      v = v if getattr(v, "dtype", None) else np.asarray(v)
+      dtype = dtypes.canonicalize_dtype(v.dtype)
+      if dtype != v.dtype:
+        v = v.astype(dtype)
+      return v
+
+    args_flat_jax = tuple(map(canonical_arg, args_flat_jax))
+    def make_tensorspec(a_jax):
+      a_tf_dtype = jax2tf_internal._to_tf_dtype(a_jax.dtype)
+      a_tf_shape = [d if core.is_constant_dim(d) else None for d in a_jax.shape]
+      return tf.TensorSpec(a_tf_shape, a_tf_dtype)
+    args_flat_sig_tf = tuple(map(make_tensorspec, args_flat_jax))
+
+    if not isinstance(output_shape_dtype, UnspecifiedOutputShapeDtype):
+      output_shape_dtype_flat, output_shape_dtype_tree = tree_util.tree_flatten(output_shape_dtype)
+      output_avals = tuple(core.ShapedArray(st.shape, st.dtype) for st in output_shape_dtype_flat)
+    else:
+      output_avals, output_shape_dtype_tree = None, None
+
+    res_treedef = None  # We'll store here the result treedef
+    res_tf_flat = None  # For error reporting
+    # The function below will be called at least once, either in eager
+    # mode during jax2tf_call_tf or in graph mode during _get_concrete_function_tf()
+    def callable_flat_tf(*args_tf_flat: TfVal) -> Sequence[TfVal]:
+      args_tf = args_treedef.unflatten(args_tf_flat)
+      res_tf = callable_tf(*args_tf)
+
+      # b/279454591: When `callable_tf` is a tf function with zero outputs, it
+      # returns a `StatefulPartitionedCall` (if the function is stateful) or
+      # `PartitionedCall` (if the function is stateless) op instead of
+      # tf.Tensors. We work around this issue by replacing the output `res_tf`
+      # with an empty list.
+
+      if isinstance(res_tf, tf.Operation):
+        assert (
+            res_tf.type == "StatefulPartitionedCall"
+            or res_tf.type == "PartitionedCall"
+        )
+        t_out = res_tf.get_attr("Tout")
+        # t_out should be an empty list.
+        assert not t_out, (
+            "The TF function returned an unexpected result, please check its"
+            f" function body. res_tf = {res_tf}"
+        )
+        res_tf = t_out
+
+      nonlocal res_treedef, res_tf_flat
+      res_tf_flat, res_treedef_now = tree_util.tree_flatten(res_tf)
+      assert res_treedef is None or res_treedef == res_treedef_now, (
+          f"Subsequent calls had different results. Previous {res_treedef} and now {res_treedef_now}")
+      res_treedef = res_treedef_now
+      if output_avals is not None:
+        if res_treedef != output_shape_dtype_tree:
+          raise ValueError(
+              "The pytree of the TensorFlow function results does not match the "
+              "pytree of the declared output_shape_dtype:\n"
+              f"results pytree: {res_treedef}\noutput_shape_dtype tree: {output_shape_dtype_tree}")
+        assert len(output_avals) == len(res_tf_flat)
+
+      checked_res_tf_flat = [
+          check_tf_result(i, r_tf, r_aval)
+          for i, (r_tf, r_aval) in enumerate(
+              zip(res_tf_flat,
+                  (output_avals
+                   if output_avals is not None
+                   else (None,) * len(res_tf_flat))))]
+      return checked_res_tf_flat
+
+    # Prepare a tf.function ahead of time, to cache the concrete functions. This
+    # won't be used in op-by-op execution mode.
+    function_flat_tf = tf.function(
+        callable_flat_tf, autograph=False, jit_compile=not call_tf_graph)
+
+    res_jax_flat = call_tf_p.bind(
+        *args_flat_jax,
+        # Carry the actual function such that op-by-op call can call in TF eager mode.
+        callable_flat_tf=callable_flat_tf,
+        function_flat_tf=function_flat_tf,
+        args_flat_sig_tf=args_flat_sig_tf,
+        output_avals=output_avals,
+        has_side_effects=has_side_effects,
+        ordered=ordered,
+        call_tf_graph=call_tf_graph,
+    )
+
+    # We must have called callable_flat_tf by nοw
+    assert res_treedef is not None
+    return res_treedef.unflatten(res_jax_flat)
+
+  # Define the fwd and bwd custom_vjp functions
+  def make_call_vjp_fwd(*args_jax):
+    # Return the primal arguments as the residual
+    return make_call(*args_jax), args_jax
+
+  def make_call_vjp_bwd(residual_jax, ct_res_jax):
+    args_jax = residual_jax  # residual is the primal argument
+
+    def tf_vjp_fun(args_tf, ct_res_tf):
+      """Invoke TF gradient."""
+
+      # TF does not like us to watch non-float vars
+      def replace_non_float(arg_tf):
+        if arg_tf.dtype.is_floating or arg_tf.dtype.is_complex:
+          return arg_tf
+        else:
+          # When watched, this will be ignored. When used in results it will
+          # result in a floating 0. gradient, which JAX will ignore (and
+          # replace it with a float0)
+          return tf.zeros((), dtype=tf.float32)
+
+      watched_args_tf = tf.nest.map_structure(replace_non_float, args_tf)
+      with tf.GradientTape(persistent=True) as tape:
+        tape.watch(watched_args_tf)
+        res = callable_tf(*args_tf)
+
+      tf.nest.assert_same_structure(res, ct_res_tf)
+      dres_darg = tape.gradient(
+          tf.nest.map_structure(replace_non_float, res),
+          sources=watched_args_tf,
+          output_gradients=ct_res_tf,
+          unconnected_gradients=tf.UnconnectedGradients.ZERO)
+
+      dres_darg = tree_util.tree_map(
+          lambda x: x if x is None else tf.convert_to_tensor(x),
+          dres_darg,
+      )
+      tf.nest.assert_same_structure(dres_darg, args_tf)
+      return dres_darg
+
+    # Use call_tf to call the VJP function
+    ct_args_jax = call_tf(tf_vjp_fun)(args_jax, ct_res_jax)
+    # We must make the float0s that JAX expects
+    def fix_float0(arg_jax, ct_arg_jax):
+      arg_dtype = dtypes.result_type(arg_jax)  # May be scalar
+      ct_arg_dtype = core.primal_dtype_to_tangent_dtype(arg_dtype)
+      if ct_arg_dtype != ct_arg_jax.dtype:
+        return ad_util.zeros_like_aval(core.ShapedArray(np.shape(arg_jax),
+                                                        ct_arg_dtype))
+      return ct_arg_jax
+
+    ct_args_jax_fixed = tree_util.tree_map(fix_float0, args_jax, ct_args_jax)
+    return ct_args_jax_fixed
+
+  make_call.defvjp(make_call_vjp_fwd, make_call_vjp_bwd)
+  return util.wraps(callable_tf)(make_call)
+
+
+def check_tf_result(idx: int, r_tf: TfVal, r_aval: core.ShapedArray | None) -> TfVal:
+  # Check that the TF function returns values of expected types. This
+  # improves error reporting, preventing hard-to-diagnose errors downstream
+  try:
+    jax2tf_internal._tfval_to_tensor_jax_dtype(r_tf)
+  except Exception as e:
+    msg = ("The called TF function returns a result that is not "
+           f"convertible to JAX: {r_tf}.")
+    raise ValueError(msg) from e
+
+  if r_aval is None:
+    return r_tf
+  # We convert to TF type, and canonicalize to 32-bit if necessary
+  r_aval_dtype_tf = jax2tf_internal._to_tf_dtype(r_aval.dtype)
+  # Checking shapes is trickier in presence of dynamic shapes. I wish we could
+  # check at runtime that the returned shape matches the declared shape. I wish
+  # that tf.ensure_shape did this, but it can only take shapes that contain None
+  # not computed shapes. However, in eager mode we should be able to resolve
+  # the declared shapes to constants and we get better checking.
+  if tf.executing_eagerly():
+    r_aval_shape_tf = jax2tf_internal._eval_shape(r_aval.shape)
+  else:
+    r_aval_shape_tf = jax2tf_internal._aval_to_tf_shape(r_aval)
+  # We do as much checking as we can here, instead of relying on tf.ensure_shape
+  # because the latter gives different errors in eager vs. compiled mode.
+  # TODO(b/279454591): This strange error is from TF. Eager function suppose
+  # return tf Val with concrete shape but not.  Here we change exception to warn
+  # and bypass it. This case need revisit on TF side.
+  try:
+    _ = len(r_tf.shape)
+  except ValueError as e:
+    msg = (
+        "The shape check test cannot be performed because the shape of the"
+        "`r_tf` tensor cannot be obtained."
+        f"r_tf = {r_tf}, r_aval = {r_aval}"
+    )
+    msg += str(e)
+    logging.warning(msg)
+    return r_tf
+  if (r_tf.dtype != r_aval_dtype_tf or
+      len(r_tf.shape) != len(r_aval_shape_tf) or
+      any(r_aval_d is not None and r_tf_d is not None and r_aval_d != r_tf_d
+          for r_tf_d, r_aval_d in zip(r_tf.shape, r_aval_shape_tf))):
+    msg = ("The shapes or dtypes returned by the TensorFlow function "
+           "do not match the declared output_shape_dtype:\n"
+           f"Result[{idx}] is {r_tf.dtype}[{r_tf.shape}] vs. expected {r_aval_dtype_tf}[{r_aval_shape_tf}]")
+    raise ValueError(msg)
+  # At this point tf.ensure_shape does not do much, it should never throw an
+  # error, albeit it may refine the shape a bit.
+  return tf.ensure_shape(r_tf, r_aval_shape_tf)
+
+
+call_tf_p = core.Primitive("call_tf")
+call_tf_p.multiple_results = True
+
+# The impl will be used in op-by-op mode and calls callable_tf in TF eager mode.
+def _call_tf_impl(*args_jax_flat, callable_flat_tf, **_):
+  # On GPU we use dlpack to avoid copies of data to the host.
+  def _arg_jax_to_tf(arg_jax):
+    if (isinstance(arg_jax, jax.Array) and
+        list(arg_jax.devices())[0].platform in _DLPACK_PLATFORMS and
+        arg_jax.dtype.type in dlpack.SUPPORTED_DTYPES):
+      arg_dlpack = jax.dlpack.to_dlpack(arg_jax)
+      return tf.experimental.dlpack.from_dlpack(arg_dlpack)
+    # The following avoids copies to the host on CPU, always for Array
+    # and even for ndarray if they are sufficiently aligned.
+    # TODO(necula): on TPU this copies to the host!
+    if getattr(arg_jax, 'dtype', None) == dtypes.float0:
+      return tf.zeros(shape=arg_jax.shape,
+                      dtype=jax2tf_internal._tf_np_dtype_for_float0)
+    return tf.constant(np.asarray(arg_jax))
+
+  args_tf_flat = tuple(map(_arg_jax_to_tf, args_jax_flat))
+  with jax2tf_internal.inside_call_tf():
+    # Call in TF eager mode
+    res_tf_flat = callable_flat_tf(*args_tf_flat)
+
+  def _res_tf_to_jax(res_tf: TfVal):
+    res_tf, jax_dtype = jax2tf_internal._tfval_to_tensor_jax_dtype(res_tf)
+    if isinstance(res_tf, tf.Tensor) and jax_dtype.type in dlpack.SUPPORTED_DTYPES:
+      res_tf_platform = tf.DeviceSpec.from_string(res_tf.backing_device).device_type
+      res_jax_platform = res_tf_platform.lower()
+      if res_jax_platform in _DLPACK_PLATFORMS:
+        res_dlpack = tf.experimental.dlpack.to_dlpack(res_tf)
+        return jax.dlpack.from_dlpack(res_dlpack)
+
+    # When working with a bfloat16 scalar tf.Tensor,np.asarray() can fail.
+    # To handle this special case, we create a numpy copy.
+    if res_tf.shape == tf.TensorShape([]) and res_tf.dtype == tf.bfloat16:
+      return jax.device_put(jnp.array(res_tf.numpy()))
+    else:
+      return jax.device_put(np.asarray(res_tf))
+
+  return list(map(_res_tf_to_jax, res_tf_flat))
+
+
+call_tf_p.def_impl(_call_tf_impl)
+
+@functools.lru_cache(maxsize=128)
+def _get_concrete_function_tf(function_flat_tf, args_flat_sig_tf):  # -> tf.ConcreteFunction
+  with jax2tf_internal.inside_call_tf():
+    return function_flat_tf.get_concrete_function(*args_flat_sig_tf)
+
+
+# Mark the effectful instances of call_tf
+@dataclasses.dataclass(frozen=True)
+class CallTfEffect(effects.Effect):
+  __str__ = lambda _: "CallTfEffect"
+
+call_tf_effect = CallTfEffect()
+
+effects.lowerable_effects.add_type(CallTfEffect)
+effects.control_flow_allowed_effects.add_type(CallTfEffect)
+effects.remat_allowed_effects.add_type(CallTfEffect)
+effects.custom_derivatives_allowed_effects.add_type(CallTfEffect)
+
+
+class CallTfOrderedEffect(effects.Effect):
+  __str__ = lambda _: "CallTfOrderedEffect"
+
+
+call_tf_ordered_effect = CallTfOrderedEffect()
+
+effects.lowerable_effects.add_type(CallTfOrderedEffect)
+effects.control_flow_allowed_effects.add_type(CallTfOrderedEffect)
+effects.remat_allowed_effects.add_type(CallTfOrderedEffect)
+effects.custom_derivatives_allowed_effects.add_type(CallTfOrderedEffect)
+effects.ordered_effects.add_type(CallTfOrderedEffect)
+effects.shardable_ordered_effects.add_type(CallTfOrderedEffect)
+
+
+def _call_tf_abstract_eval(
+    *args_flat_avals,
+    function_flat_tf,
+    args_flat_sig_tf,
+    has_side_effects,
+    ordered,
+    output_avals,
+    call_tf_graph,
+    **__,
+):
+  # Called only when we form a Jaxpr, i.e., under jit, scan, etc.
+  effects = set()
+  if ordered:
+    effects.add(call_tf_ordered_effect)
+  elif has_side_effects:
+    effects.add(call_tf_effect)
+
+  # If no output_avals is given, then we ask TF to infer the output shapes.
+  # We call this even if output_avals is given because it will ensure that
+  # callable_flat_tf is called. Since _get_concrete_function_tf is cached
+  # there is a small cost of calling it more often than needed.
+  concrete_function_flat_tf = _get_concrete_function_tf(function_flat_tf,
+                                                        args_flat_sig_tf)
+
+  # In the case that the tf.function has no return value
+  if len(concrete_function_flat_tf.outputs) == 0:
+    return (), effects
+
+  if output_avals is not None:
+    return output_avals, effects
+
+  def is_fully_known_shape(s):
+    return s.rank is not None and all(d is not None for d in s)
+
+  if all(is_fully_known_shape(s)
+        for s in concrete_function_flat_tf.output_shapes):
+    avals_from_tf = tuple(
+        # We convert to JAX type, and canonicalize to 32-bit if necessary
+        core.ShapedArray(shape, jax2tf_internal._to_jax_dtype(dtype))
+        for dtype, shape in zip(concrete_function_flat_tf.output_dtypes,
+                                concrete_function_flat_tf.output_shapes))
+    return avals_from_tf, effects
+
+  msg = ("call_tf cannot call functions whose output has dynamic shape. "
+    f"Found output shapes: {concrete_function_flat_tf.output_shapes}. "
+    "Consider using the `output_shape_dtype` argument to call_tf. "
+    "\nSee https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#limitations-of-call_tf"
+      " for a discussion.")
+  raise ValueError(msg)
+
+
+call_tf_p.def_effectful_abstract_eval(_call_tf_abstract_eval)
+
+
+def _call_tf_lowering(
+    ctx: mlir.LoweringRuleContext,
+    *args_op,
+    platform,
+    function_flat_tf,
+    args_flat_sig_tf,
+    has_side_effects,
+    ordered,
+    call_tf_graph,
+    output_avals,
+    **_,
+):
+  # We use the same TF lowering device as for the embedding JAX computation.
+  # One example when this is needed is when the code refers to variables on one
+  # device. Or, for sharding annotations (only supported on TPU).
+
+  if platform in ["cpu", "tpu"]:
+    tf_platform = platform.upper()
+  elif platform == "cuda":
+    tf_platform = "GPU"
+  else:
+    raise ValueError("platform {platform} not supported")
+
+  concrete_function_flat_tf = _get_concrete_function_tf(function_flat_tf, args_flat_sig_tf)
+
+  captured_inputs = []
+  if concrete_function_flat_tf.captured_inputs:
+    # The function uses either captured variables or tensors.
+    msg = (
+        "call_tf works best with a TensorFlow function that does not capture "
+        "variables or tensors from the context. "
+        "See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#limitations-of-call_tf for a discussion. "
+        f"The following captures were found {concrete_function_flat_tf.captured_inputs}")
+    logging.warning(msg)
+    for inp in concrete_function_flat_tf.captured_inputs:
+      if inp.dtype == tf.resource:  # A variable; lookup by handle
+        inp_vars = [v for v in concrete_function_flat_tf.variables if inp is v.handle]
+        assert len(inp_vars) == 1, f"Found {inp_vars}"
+        captured_inputs.append(inp_vars[0])
+      else:
+        captured_inputs.append(inp)
+
+  captured_ops = tuple(
+      mlir.ir_constant(np.asarray(inp))
+      for inp in captured_inputs
+  )
+
+  if call_tf_graph:
+    with jax2tf_internal.inside_call_tf():
+      return emit_tf_embedded_graph_custom_call(
+          ctx,
+          concrete_function_flat_tf,
+          tuple(args_op) + captured_ops,
+          has_side_effects,
+          ordered,
+          output_avals,
+      )
+
+  def convert_to_spec(x):
+    if isinstance(x, tf.TensorSpec):
+      return x
+    else:
+      return tf.TensorSpec.from_tensor(x)
+
+  args_tf_flat = [convert_to_spec(a) for a in args_flat_sig_tf]
+
+  with jax2tf_internal.inside_call_tf():
+    try:
+      func_tf_hlo = function_flat_tf.experimental_get_compiler_ir(
+          *args_tf_flat
+      )(stage="hlo_serialized", platform_name=tf_platform)
+    except Exception as e:
+      msg = ("Error compiling TensorFlow function (see below for the caught exception)." +
+             "\ncall_tf can used " +
+              "in a staged context (under jax.jit, lax.scan, etc.) only with " +
+              "compilable functions with static output shapes.\n" +
+              "See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#limitations-of-call_tf for a discussion." +
+             "\n\nCaught TensorFlow exception: " + str(e))
+      raise ValueError(msg) from e
+
+  xla_comp = xla_client.XlaComputation(func_tf_hlo)
+
+  # Canonicalize the results; e.g., makes them x32 if JAX is in 32-bit mode
+  def canonical_res_aval(res_shape: xla_client.Shape) -> core.ShapedArray:
+    if not res_shape.is_static():
+      msg = ("Compiled TensorFlow function has dynamic output shape " +
+             f"{res_shape}. call_tf can used " +
+             "in a staged context (under jax.jit, lax.scan, etc.) only with " +
+             "compilable functions with static output shapes. " +
+             "See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#limitations-of-call_tf for a discussion.")
+      raise ValueError(msg)
+
+    res_dtype = res_shape.numpy_dtype()
+    jax_res_dtype = dtypes.canonicalize_dtype(res_dtype)
+    return core.ShapedArray(res_shape.dimensions(), jax_res_dtype)
+
+  result_shape = xla_comp.program_shape().result_shape()
+  if not result_shape.is_tuple():
+    # TF does not wrap singletons as tuples, but JAX expects tuples because
+    # call_tf is a multiple_results primitive.
+    result_shapes = (result_shape,)
+  else:
+    result_shapes = result_shape.tuple_shapes()  # type: ignore
+
+  result_avals = tuple(map(canonical_res_aval, result_shapes))
+
+  submodule = mlir.xla_computation_to_mlir_module(xla_comp)
+  symtab = ir.SymbolTable(submodule.operation)
+  callee_result_types = symtab["main"].type.results
+  fn = mlir.merge_mlir_modules(ctx.module_context.module,
+                               f"call_tf_{function_flat_tf.name}",
+                               submodule,
+                               dst_symtab=ctx.module_context.symbol_table)
+  call = func_dialect.CallOp(callee_result_types,
+                             ir.FlatSymbolRefAttr.get(fn),
+                             tuple(args_op) + captured_ops)
+  if result_shape.is_tuple():
+    flat_results = [hlo.get_tuple_element(call, mlir.i32_attr(i))
+                    for i in range(len(result_shapes))]
+  else:
+    flat_results = call.results
+
+  if ordered:
+    raise NotImplementedError(
+        "ordered=True is not supported in the jitted context without"
+        " `call_tf_graph=True`"
+    )
+
+  outputs = []
+  for op, res_aval, res_shape in zip(flat_results, result_avals,
+                                     result_shapes):
+    if res_aval.dtype != res_shape.numpy_dtype():
+      op = hlo.ConvertOp(mlir.aval_to_ir_type(res_aval), op).result
+    outputs.append(op)
+  return outputs
+
+
+def _register_call_lowering(platform):
+  mlir.register_lowering(call_tf_p, functools.partial(_call_tf_lowering,
+                                                      platform=platform),
+                         platform=platform)
+for platform in ("cpu", "cuda", "tpu"):
+  _register_call_lowering(platform)
+
+# Support the call_tf under jax2tf.convert in eager mode
+def _jax2tf_call_tf(*args: TfVal,
+                    callable_flat_tf: Callable,
+                    **_) -> TfVal:
+  with jax2tf_internal.inside_call_tf():
+    res_tf_flat = callable_flat_tf(*args)
+  return res_tf_flat
+
+jax2tf_internal.tf_impl[call_tf_p] = _jax2tf_call_tf
+
+
+def emit_tf_embedded_graph_custom_call(
+    ctx: mlir.LoweringRuleContext,
+    concrete_function_flat_tf,
+    operands: Sequence[ir.Value],
+    has_side_effects,
+    ordered,
+    output_avals,
+):
+  """Emits a custom call referencing a tf.Graph embedding of the TF function.
+
+  All call_tf called function information is stored in tf.metadata.
+  This includes:
+  (1) The called function name: This name will be used by the runtime to execute
+  the callback.
+  (2) The called function index in the XLACallModule `function_list` attribute.
+  """
+  call_tf_concrete_function_list = jax2tf_internal.get_thread_local_state_call_tf_concrete_function_list()
+  if call_tf_concrete_function_list is None:
+    raise ValueError(
+        "call_tf_graph=True only support exporting by jax2tf.convert currently."
+    )
+  # TODO(necula): It is dangerous to modify global state when lowering because
+  # there are a number of lowering caches that only cache the StableHLO.
+  # See call_tf_test.py:test_multi_platform_call_tf_graph.
+  called_index = add_to_call_tf_concrete_function_list(
+      concrete_function_flat_tf, call_tf_concrete_function_list)
+  tf_backend_config = {
+      "has_token_input_output": ir.BoolAttr.get(ordered),
+      "called_index": mlir.i64_attr(called_index),
+  }
+  result_avals = ctx.avals_out if ctx.avals_out is not None else ()
+
+  operands = list(operands)
+  result_types = list(
+      util.flatten([mlir.aval_to_ir_types(aval) for aval in result_avals])
+  )
+  if ordered:
+    operands.insert(0, ctx.tokens_in.get(call_tf_ordered_effect)[0])
+    result_types.insert(0, mlir.token_type()[0])
+
+  custom_call = hlo.CustomCallOp(
+      result_types,
+      operands,
+      call_target_name=ir.StringAttr.get("tf.call_tf_function"),
+      has_side_effect=ir.BoolAttr.get(has_side_effects),
+      api_version=mlir.i32_attr(2),
+      called_computations=ir.ArrayAttr.get([]),
+      backend_config=ir.StringAttr.get(""),
+  )
+  # Store TF metadata in unregistered attribute
+  custom_call.attributes["tf.backend_config"] = ir.DictAttr.get(
+      tf_backend_config
+  )
+
+  results = list(custom_call.results)
+  if ordered:
+    token = results.pop(0)
+    ctx.set_tokens_out(mlir.TokenSet({call_tf_ordered_effect: (token,)}))
+
+  return results
+
+
+def add_to_call_tf_concrete_function_list(concrete_tf_fn: Any, call_tf_concrete_function_list: list[Any]) -> int:
+  try:
+    called_index = call_tf_concrete_function_list.index(concrete_tf_fn)
+  except ValueError:
+    called_index = len(call_tf_concrete_function_list)
+    call_tf_concrete_function_list.append(concrete_tf_fn)
+  return called_index

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/examples/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2021 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/examples/keras_reuse_main.py

@@ -0,0 +1,78 @@
+# Copyright 2020 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Demonstrates reuse of a jax2tf model in Keras.
+
+Includes the flags from saved_model_main.py.
+
+See README.md.
+"""
+import logging
+from absl import app
+from absl import flags
+from jax.experimental.jax2tf.examples import mnist_lib
+from jax.experimental.jax2tf.examples import saved_model_main
+import tensorflow as tf
+import tensorflow_datasets as tfds  # type: ignore
+import tensorflow_hub as hub  # type: ignore
+
+
+FLAGS = flags.FLAGS
+
+
+def main(_):
+  FLAGS.model_classifier_layer = False  # We only need the features
+  # Train the model and save the feature extractor
+  saved_model_main.train_and_save()
+
+  tf_accelerator, _ = saved_model_main.tf_accelerator_and_tolerances()
+  feature_model_dir = saved_model_main.savedmodel_dir()
+
+  # With Keras, we use the tf.distribute.OneDeviceStrategy as the high-level
+  # analogue of the tf.device(...) placement seen above.
+  # It works on CPU, GPU and TPU.
+  # Actual high-performance training would use the appropriately replicated
+  # TF Distribution Strategy.
+  strategy = tf.distribute.OneDeviceStrategy(tf_accelerator)
+  with strategy.scope():
+    images = tf.keras.layers.Input(
+        mnist_lib.input_shape, batch_size=mnist_lib.train_batch_size)
+    keras_feature_extractor = hub.KerasLayer(feature_model_dir, trainable=True)
+    features = keras_feature_extractor(images)
+    predictor = tf.keras.layers.Dense(10, activation="softmax")
+    predictions = predictor(features)
+    keras_model = tf.keras.Model(images, predictions)
+
+  keras_model.compile(
+      loss=tf.keras.losses.categorical_crossentropy,
+      optimizer=tf.keras.optimizers.SGD(learning_rate=0.01),
+      metrics=["accuracy"])
+  logging.info(keras_model.summary())
+
+  train_ds = mnist_lib.load_mnist(
+      tfds.Split.TRAIN, batch_size=mnist_lib.train_batch_size)
+  test_ds = mnist_lib.load_mnist(
+      tfds.Split.TEST, batch_size=mnist_lib.test_batch_size)
+  keras_model.fit(train_ds, epochs=FLAGS.num_epochs, validation_data=test_ds)
+
+  if saved_model_main.SHOW_IMAGES.value:
+    mnist_lib.plot_images(
+        test_ds,
+        1,
+        5,
+        f"Keras inference with reuse of {saved_model_main.model_description()}",
+        inference_fn=lambda images: keras_model(tf.convert_to_tensor(images)))
+
+
+if __name__ == "__main__":
+  app.run(main)

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/examples/keras_reuse_main_test.py

@@ -0,0 +1,50 @@
+# Copyright 2020 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from absl import flags
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+from jax._src import test_util as jtu
+
+from jax.experimental.jax2tf.examples import keras_reuse_main
+from jax.experimental.jax2tf.tests import tf_test_util
+
+jax.config.parse_flags_with_absl()
+FLAGS = flags.FLAGS
+
+
+class KerasReuseMainTest(tf_test_util.JaxToTfTestCase):
+
+  def setUp(self):
+    super().setUp()
+    FLAGS.model_path = os.path.join(absltest.get_default_test_tmpdir(),
+                                    "saved_models")
+    FLAGS.num_epochs = 1
+    FLAGS.test_savedmodel = True
+    FLAGS.mock_data = True
+    FLAGS.show_images = False
+    FLAGS.serving_batch_size = 1
+
+  @parameterized.named_parameters(
+      dict(testcase_name=f"_{model}", model=model)
+      for model in ["mnist_pure_jax", "mnist_flax"])
+  def test_keras_reuse(self, model="mnist_pure_jax"):
+    FLAGS.model = model
+    keras_reuse_main.main(None)
+
+
+if __name__ == "__main__":
+  absltest.main(testLoader=jtu.JaxTestLoader())

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/examples/mnist_lib.py

@@ -0,0 +1,324 @@
+# Copyright 2020 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Definitions of two versions of MNIST (model and training code ).
+
+One definition uses pure JAX (for those who prefer an example with fewer
+moving parts, at the expense of code size), and another using Flax.
+
+See README.md for how these are used.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+import functools
+import logging
+import re
+import time
+from typing import Any, Callable, Optional
+from absl import flags
+
+import flax
+from flax import linen as nn
+
+import jax
+import jax.numpy as jnp
+
+from matplotlib import pyplot as plt
+import numpy as np
+import optax
+import tensorflow as tf
+import tensorflow_datasets as tfds  # type: ignore
+
+_MOCK_DATA = flags.DEFINE_boolean("mock_data", False,
+                                  "Use fake data, for testing.")
+
+#### Model parameters
+
+# For fun, let's use different batch sizes for training and for evaluation.
+train_batch_size = 128
+test_batch_size = 16
+
+# Define common parameters for both the JAX and the Flax models.
+input_shape = (28, 28, 1)  # Excluding batch_size
+layer_sizes = [784, 512, 512, 10]  # 10 is the number of classes
+param_scale = 0.1
+step_size = 0.001
+
+
+def load_mnist(split: tfds.Split, batch_size: int):
+  """Loads either training or test MNIST data.
+
+  Args:
+    split: either tfds.Split.TRAIN or tfds.Split.TEST.
+
+  Returns:
+    an iterator with pairs (images, labels). The images have shape
+    (B, 28, 28, 1) and the labels have shape (B, 10), where B is the batch_size.
+  """
+  if _MOCK_DATA.value:
+    with tfds.testing.mock_data(num_examples=batch_size):
+      try:
+        ds = tfds.load("mnist", split=split)
+      except Exception as e:
+        m = re.search(r'metadata files were not found in (.+/)mnist/', str(e))
+        if m:
+          msg = ("TFDS mock_data is missing the mnist metadata files. Run the "
+                 "`saved_model_main.py` binary and see where TFDS downloads "
+                 "the mnist data set (typically ~/tensorflow_datasets/mnist). "
+                 f"Copy the `mnist` directory to {m.group(1)} and re-run the test")
+          raise ValueError(msg) from e
+        else:
+          raise e
+  else:
+    ds = tfds.load("mnist", split=split)
+
+  def _prepare_example(x):
+    image = tf.cast(x["image"], tf.float32) / 255.0
+    label = tf.one_hot(x["label"], 10)
+    return (image, label)
+
+  ds = ds.map(_prepare_example)
+  # drop_remainder=True is important for use with Keras
+  ds = ds.cache().shuffle(1000).batch(batch_size, drop_remainder=True)
+  return ds
+
+
+class PureJaxMNIST:
+  """An MNIST model written using pure JAX.
+
+  There is an option for the model to skip the classifier layer, for
+  demonstrating reuse of the classifier-less model into a larger model.
+  See README.md.
+  """
+
+  name = "mnist_pure_jax"
+
+  @staticmethod
+  def predict(params: Sequence[tuple[Any, Any]], inputs, with_classifier=True):
+    """The prediction function.
+
+    Args:
+      params: a list with pairs of weights and biases for each layer.
+      inputs: the batch of images (B, 28, 28, 1)
+      with_classifier: whether to include the classifier layer.
+
+    Returns:
+      either the predictions (B, 10) if with_classifier=True, or the
+      final set of logits of shape (B, 512).
+    """
+    x = inputs.reshape((inputs.shape[0], -1))  # flatten to f32[B, 784]
+    for w, b in params[:-1]:
+      x = jnp.dot(x, w) + b
+      x = jnp.tanh(x)
+
+    if not with_classifier:
+      return x
+    final_w, final_b = params[-1]
+    logits = jnp.dot(x, final_w) + final_b
+    return logits - jax.scipy.special.logsumexp(
+      logits, axis=1, keepdims=True)
+
+  @staticmethod
+  def loss(params, inputs, labels):
+    predictions = PureJaxMNIST.predict(params, inputs, with_classifier=True)
+    return -jnp.mean(jnp.sum(predictions * labels, axis=1))
+
+  @staticmethod
+  def accuracy(predict: Callable, params, dataset):
+
+    @jax.jit
+    def _per_batch(inputs, labels):
+      target_class = jnp.argmax(labels, axis=1)
+      predicted_class = jnp.argmax(predict(params, inputs), axis=1)
+      return jnp.mean(predicted_class == target_class)
+
+    batched = [
+      _per_batch(inputs, labels) for inputs, labels in tfds.as_numpy(dataset)
+    ]
+    return jnp.mean(jnp.stack(batched))
+
+  @staticmethod
+  def update(params, inputs, labels):
+    grads = jax.grad(PureJaxMNIST.loss)(params, inputs, labels)
+    return [(w - step_size * dw, b - step_size * db)
+            for (w, b), (dw, db) in zip(params, grads)]
+
+  @staticmethod
+  def train(train_ds, test_ds, num_epochs, with_classifier=True):
+    """Trains a pure JAX MNIST predictor.
+
+    Returns:
+      a tuple with two elements:
+        - a predictor function with signature "(Params, ImagesBatch) ->
+        Predictions".
+          If `with_classifier=False` then the output of the predictor function
+          is the last layer of logits.
+        - the parameters "Params" for the predictor function
+    """
+    rng = jax.random.PRNGKey(0)
+    params = [(param_scale * jax.random.normal(rng, (m, n)),
+               param_scale * jax.random.normal(rng, (n,)))
+              for m, n, in zip(layer_sizes[:-1], layer_sizes[1:])]
+
+    for epoch in range(num_epochs):
+      start_time = time.time()
+      for inputs, labels in tfds.as_numpy(train_ds):
+        params = jax.jit(PureJaxMNIST.update)(params, inputs, labels)
+      epoch_time = time.time() - start_time
+      train_acc = PureJaxMNIST.accuracy(PureJaxMNIST.predict, params, train_ds)
+      test_acc = PureJaxMNIST.accuracy(PureJaxMNIST.predict, params, test_ds)
+      logging.info("%s: Epoch %d in %0.2f sec", PureJaxMNIST.name, epoch,
+                   epoch_time)
+      logging.info("%s: Training set accuracy %0.2f%%", PureJaxMNIST.name,
+                   100. * train_acc)
+      logging.info("%s: Test set accuracy %0.2f%%", PureJaxMNIST.name,
+                   100. * test_acc)
+
+    return (functools.partial(
+      PureJaxMNIST.predict, with_classifier=with_classifier), params)
+
+
+class FlaxMNIST:
+  """An MNIST model using Flax."""
+
+  name = "mnist_flax"
+
+  class Module(nn.Module):
+    """A simple CNN model for MNIST.
+
+    There is an option for the model to skip the classifier layer, for
+    demonstrating reuse of the classifier-less model into a larger model.
+    See README.md.
+    """
+
+    @nn.compact
+    def __call__(self, x, with_classifier=True):
+      x = nn.Conv(features=32, kernel_size=(3, 3))(x)
+      x = nn.relu(x)
+      x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
+      x = nn.Conv(features=64, kernel_size=(3, 3))(x)
+      x = nn.relu(x)
+      x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
+      x = x.reshape((x.shape[0], -1))  # flatten
+      x = nn.Dense(features=256)(x)
+      x = nn.relu(x)
+      if not with_classifier:
+        return x
+      x = nn.Dense(features=10)(x)
+      x = nn.log_softmax(x)
+      return x
+
+  # Create the model and save it
+  model = Module()
+
+  @staticmethod
+  def predict(params, inputs, with_classifier=True):
+    return FlaxMNIST.model.apply({"params": params},
+                                 inputs,
+                                 with_classifier=with_classifier)
+
+  @staticmethod
+  def loss(params, inputs, labels):  # Same as the pure JAX example
+    # Must use the classifier layer because the labels are classes
+    predictions = FlaxMNIST.predict(params, inputs, with_classifier=True)
+    return -jnp.mean(jnp.sum(predictions * labels, axis=1))
+
+  @staticmethod
+  def update(tx, params, opt_state, inputs, labels):
+    grad = jax.grad(FlaxMNIST.loss)(params, inputs, labels)
+    updates, opt_state = tx.update(grad, opt_state)
+    params = optax.apply_updates(params, updates)
+    return params, opt_state
+
+  @staticmethod
+  def train(train_ds, test_ds, num_epochs, with_classifier=True):
+    """Trains a pure JAX MNIST predictor.
+
+    Returns:
+      a tuple with two elements:
+        - a predictor function with signature "(Params, ImagesBatch) ->
+          Predictions".
+          If `with_classifier=False` then the output of the predictor function
+          is the last layer of logits.
+        - the parameters "Params" for the predictor function
+    """
+    rng = jax.random.PRNGKey(0)
+    momentum_mass = 0.9
+
+    init_shape = jnp.ones((1,) + input_shape, jnp.float32)
+    params = FlaxMNIST.model.init(rng, init_shape)["params"]
+    tx = optax.sgd(learning_rate=step_size, momentum=momentum_mass)
+    opt_state = tx.init(params)
+
+    for epoch in range(num_epochs):
+      start_time = time.time()
+      for inputs, labels in tfds.as_numpy(train_ds):
+        params, opt_state = jax.jit(FlaxMNIST.update,
+                                    static_argnums=0)(tx, params, opt_state,
+                                                      inputs, labels)
+      epoch_time = time.time() - start_time
+      # Same accuracy function as for the pure JAX example
+      train_acc = PureJaxMNIST.accuracy(FlaxMNIST.predict, params,
+                                        train_ds)
+      test_acc = PureJaxMNIST.accuracy(FlaxMNIST.predict, params,
+                                       test_ds)
+      logging.info("%s: Epoch %d in %0.2f sec", FlaxMNIST.name, epoch,
+                   epoch_time)
+      logging.info("%s: Training set accuracy %0.2f%%", FlaxMNIST.name,
+                   100. * train_acc)
+      logging.info("%s: Test set accuracy %0.2f%%", FlaxMNIST.name,
+                   100. * test_acc)
+
+    # See discussion in README.md for packaging Flax models for conversion
+    predict_fn = functools.partial(FlaxMNIST.predict,
+                                   with_classifier=with_classifier)
+    return (predict_fn, params)
+
+
+def plot_images(ds,
+                nr_rows: int,
+                nr_cols: int,
+                title: str,
+                inference_fn: Callable | None = None):
+  """Plots a grid of images with their predictions.
+
+  Params:
+    ds: a tensorflow dataset from where to pick the images and labels.
+    nr_rows, nr_cols: the size of the grid to plot
+    title: the title of the plot
+    inference_fn: if None then print the existing label, else use this function
+      on the batch of images to produce a batch of inference results, which
+      get printed.
+    inference_batch_size: the size of the batch of images passed to
+    `inference_fn`.
+  """
+  count = nr_rows * nr_cols
+  fig = plt.figure(figsize=(8., 4.), num=title)
+  # Get the first batch
+  (images, labels), = list(tfds.as_numpy(ds.take(1)))
+  if inference_fn:
+    inferred_labels = inference_fn(images)
+  for i, image in enumerate(images[:count]):
+    digit = fig.add_subplot(nr_rows, nr_cols, i + 1)
+    if inference_fn:
+      digit_title = f"infer: {np.argmax(inferred_labels[i])}\n"
+    else:
+      digit_title = ""
+    digit_title += f"label: {np.argmax(labels[i])}"
+    digit.set_title(digit_title)
+    plt.imshow(
+      (np.reshape(image, (28, 28)) * 255).astype(np.uint8),
+      interpolation="nearest")
+  plt.show()

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/examples/saved_model_lib.py

@@ -0,0 +1,154 @@
+# Copyright 2020 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Defines a helper function for creating a SavedModel from a jax2tf trained model.
+
+This has been tested with TensorFlow Hub, TensorFlow JavaScript,
+and TensorFlow Serving.
+
+Note that the code in this file is provided only as an example. The functions
+generated by `jax2tf.convert` are standard TensorFlow functions and you can
+save them in a SavedModel using standard TensorFlow code. This decoupling
+of jax2tf from SavedModel is important, because it allows the user to have full
+control over what metadata is saved in the SavedModel. Please copy and
+customize this function as needed.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any, Callable, Optional, Union
+
+from jax.experimental import jax2tf
+import tensorflow as tf
+
+
+def convert_and_save_model(
+    jax_fn: Callable[[Any, Any], Any],
+    params,
+    model_dir: str,
+    *,
+    input_signatures: Sequence[tf.TensorSpec],
+    polymorphic_shapes: str | None = None,
+    with_gradient: bool = False,
+    enable_xla: bool = True,
+    compile_model: bool = True,
+    saved_model_options: tf.saved_model.SaveOptions | None = None):
+  """Convert a JAX function and saves a SavedModel.
+
+  This is an example, we do not promise backwards compatibility for this code.
+  For serious uses, please copy and expand it as needed (see note at the top
+  of the module).
+
+  Use this function if you have a trained ML model that has both a prediction
+  function and trained parameters, which you want to save separately from the
+  function graph as variables (e.g., to avoid limits on the size of the
+  GraphDef, or to enable fine-tuning.) If you don't have such parameters,
+  you can still use this library function but probably don't need it
+  (see jax2tf/README.md for some simple examples).
+
+  In order to use this wrapper you must first convert your model to a function
+  with two arguments: the parameters and the input on which you want to do
+  inference. Both arguments may be np.ndarray or (nested)
+  tuples/lists/dictionaries thereof.
+
+  See the README.md for a discussion of how to prepare Flax and Haiku models.
+
+  Args:
+    jax_fn: a JAX function taking two arguments, the parameters and the inputs.
+      Both arguments may be (nested) tuples/lists/dictionaries of np.ndarray.
+    params: the parameters, to be used as first argument for `jax_fn`. These
+      must be (nested) tuples/lists/dictionaries of np.ndarray, and will be
+      saved as the variables of the SavedModel.
+    model_dir: the directory where the model should be saved.
+    input_signatures: the input signatures for the second argument of `jax_fn`
+      (the input). A signature must be a `tensorflow.TensorSpec` instance, or a
+      (nested) tuple/list/dictionary thereof with a structure matching the
+      second argument of `jax_fn`. The first input_signature will be saved as
+      the default serving signature. The additional signatures will be used
+      only to ensure that the `jax_fn` is traced and converted to TF for the
+      corresponding input shapes.
+    with_gradient: the value to use for the `with_gradient` parameter for
+      `jax2tf.convert`.
+    enable_xla: the value to use for the `enable_xla` parameter for
+      `jax2tf.convert`.
+    compile_model: use TensorFlow jit_compiler on the SavedModel. This
+      is needed if the SavedModel will be used for TensorFlow serving.
+    polymorphic_shapes: if given then it will be used as the
+      `polymorphic_shapes` argument to jax2tf.convert for the second parameter of
+      `jax_fn`. In this case, a single `input_signatures` is supported, and
+      should have `None` in the polymorphic dimensions.
+    saved_model_options: options to pass to savedmodel.save.
+  """
+  if not input_signatures:
+    raise ValueError("At least one input_signature must be given")
+  if polymorphic_shapes is not None:
+    if len(input_signatures) > 1:
+      raise ValueError("For shape-polymorphic conversion a single "
+                       "input_signature is supported.")
+  tf_fn = jax2tf.convert(
+    jax_fn,
+    with_gradient=with_gradient,
+    polymorphic_shapes=[None, polymorphic_shapes],
+    enable_xla=enable_xla)
+
+  # Create tf.Variables for the parameters. If you want more useful variable
+  # names, you can use `tree.map_structure_with_path` from the `dm-tree` package
+  param_vars = tf.nest.map_structure(
+    lambda param: tf.Variable(param, trainable=with_gradient),
+    params)
+  tf_graph = tf.function(lambda inputs: tf_fn(param_vars, inputs),
+                         autograph=False,
+                         jit_compile=compile_model)
+
+  signatures = {}
+  # This signature is needed for TensorFlow Serving use.
+  signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \
+    tf_graph.get_concrete_function(input_signatures[0])
+  for input_signature in input_signatures[1:]:
+    # If there are more signatures, trace and cache a TF function for each one
+    tf_graph.get_concrete_function(input_signature)
+  wrapper = _ReusableSavedModelWrapper(tf_graph, param_vars)
+  if with_gradient:
+    if not saved_model_options:
+      saved_model_options = tf.saved_model.SaveOptions(experimental_custom_gradients=True)
+    else:
+      saved_model_options.experimental_custom_gradients = True
+  tf.saved_model.save(wrapper, model_dir, signatures=signatures,
+                      options=saved_model_options)
+
+
+class _ReusableSavedModelWrapper(tf.train.Checkpoint):
+  """Wraps a function and its parameters for saving to a SavedModel.
+
+  Implements the interface described at
+  https://www.tensorflow.org/hub/reusable_saved_models.
+  """
+
+  def __init__(self, tf_graph, param_vars):
+    """Args:
+
+      tf_graph: a tf.function taking one argument (the inputs), which can be
+         be tuples/lists/dictionaries of np.ndarray or tensors. The function
+         may have references to the tf.Variables in `param_vars`.
+      param_vars: the parameters, as tuples/lists/dictionaries of tf.Variable,
+         to be saved as the variables of the SavedModel.
+    """
+    super().__init__()
+    # Implement the interface from https://www.tensorflow.org/hub/reusable_saved_models
+    self.variables = tf.nest.flatten(param_vars)
+    self.trainable_variables = [v for v in self.variables if v.trainable]
+    # If you intend to prescribe regularization terms for users of the model,
+    # add them as @tf.functions with no inputs to this list. Else drop this.
+    self.regularization_losses = []
+    self.__call__ = tf_graph

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/examples/saved_model_main.py

@@ -0,0 +1,210 @@
+# Copyright 2020 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Demonstrates training models and saving the result as a SavedModel.
+
+By default, uses a pure JAX implementation of MNIST. There are flags to choose
+a Flax CNN version of MNIST, or to skip the training and just test a
+previously saved SavedModel. It is possible to save a batch-polymorphic
+version of the model, or a model prepared for specific batch sizes.
+
+Try --help to see all flags.
+
+This file is used both as an executable, and as a library in two other examples.
+See discussion in README.md.
+"""
+
+import logging
+import os
+
+from absl import app
+from absl import flags
+
+from jax.experimental.jax2tf.examples import mnist_lib
+from jax.experimental.jax2tf.examples import saved_model_lib
+
+import numpy as np
+import tensorflow as tf
+import tensorflow_datasets as tfds  # type: ignore
+
+_MODEL = flags.DEFINE_enum(
+    "model", "mnist_flax", ["mnist_flax", "mnist_pure_jax"],
+    "Which model to use.")
+_MODEL_CLASSIFIER_LAYER = flags.DEFINE_boolean("model_classifier_layer", True,
+                     ("The model should include the classifier layer, or just "
+                      "the last layer of logits. Set this to False when you "
+                      "want to reuse the classifier-less model in a larger "
+                      "model. See keras_reuse_main.py and README.md."))
+_MODEL_PATH = flags.DEFINE_string("model_path", "/tmp/jax2tf/saved_models",
+                    "Path under which to save the SavedModel.")
+_MODEL_VERSION = flags.DEFINE_integer("model_version", 1,
+                     ("The version number for the SavedModel. Needed for "
+                      "serving, larger versions will take precedence"),
+                     lower_bound=1)
+_SERVING_BATCH_SIZE = flags.DEFINE_integer("serving_batch_size", 1,
+                     "For what batch size to prepare the serving signature. "
+                     "Use -1 for converting and saving with batch polymorphism.")
+flags.register_validator(
+    "serving_batch_size",
+    lambda serving_batch_size: serving_batch_size > 0
+    or serving_batch_size == -1,
+    message="--serving_batch_size must be either -1 or a positive integer.",
+)
+
+_NUM_EPOCHS = flags.DEFINE_integer("num_epochs", 3,
+                                   "For how many epochs to train.",
+                                   lower_bound=1)
+_GENERATE_MODEL = flags.DEFINE_boolean(
+    "generate_model", True,
+    "Train and save a new model. Otherwise, use an existing SavedModel.")
+_COMPILE_MODEL = flags.DEFINE_boolean(
+    "compile_model", True,
+    "Enable TensorFlow jit_compiler for the SavedModel. This is "
+    "necessary if you want to use the model for TensorFlow serving.")
+_SHOW_MODEL = flags.DEFINE_boolean("show_model", True,
+                                   "Show details of saved SavedModel.")
+SHOW_IMAGES = flags.DEFINE_boolean(
+    "show_images", False,
+    "Plot some sample images with labels and inference results.")
+_TEST_SAVEDMODEL = flags.DEFINE_boolean(
+    "test_savedmodel", True,
+    "Test TensorFlow inference using the SavedModel w.r.t. the JAX model.")
+
+
+def train_and_save():
+  logging.info("Loading the MNIST TensorFlow dataset")
+  train_ds = mnist_lib.load_mnist(
+      tfds.Split.TRAIN, batch_size=mnist_lib.train_batch_size)
+  test_ds = mnist_lib.load_mnist(
+      tfds.Split.TEST, batch_size=mnist_lib.test_batch_size)
+
+  if SHOW_IMAGES.value:
+    mnist_lib.plot_images(train_ds, 1, 5, "Training images", inference_fn=None)
+
+  the_model_class = pick_model_class()
+  model_dir = savedmodel_dir(with_version=True)
+
+  if _GENERATE_MODEL.value:
+    model_descr = model_description()
+    logging.info("Generating model for %s", model_descr)
+    (predict_fn, predict_params) = the_model_class.train(
+        train_ds,
+        test_ds,
+        num_epochs=_NUM_EPOCHS.value,
+        with_classifier=_MODEL_CLASSIFIER_LAYER.value)
+
+    if _SERVING_BATCH_SIZE.value == -1:
+      # Batch-polymorphic SavedModel
+      input_signatures = [
+          tf.TensorSpec((None,) + mnist_lib.input_shape, tf.float32),
+      ]
+      polymorphic_shapes = "(batch, ...)"
+    else:
+      input_signatures = [
+          # The first one will be the serving signature
+          tf.TensorSpec((_SERVING_BATCH_SIZE.value,) + mnist_lib.input_shape,
+                        tf.float32),
+          tf.TensorSpec((mnist_lib.train_batch_size,) + mnist_lib.input_shape,
+                        tf.float32),
+          tf.TensorSpec((mnist_lib.test_batch_size,) + mnist_lib.input_shape,
+                        tf.float32),
+      ]
+      polymorphic_shapes = None
+
+    logging.info("Saving model for %s", model_descr)
+    saved_model_lib.convert_and_save_model(
+        predict_fn,
+        predict_params,
+        model_dir,
+        with_gradient=True,
+        input_signatures=input_signatures,
+        polymorphic_shapes=polymorphic_shapes,
+        compile_model=_COMPILE_MODEL.value)
+
+    if _TEST_SAVEDMODEL.value:
+      tf_accelerator, tolerances = tf_accelerator_and_tolerances()
+      with tf.device(tf_accelerator):
+        logging.info("Testing savedmodel")
+        pure_restored_model = tf.saved_model.load(model_dir)
+
+        if SHOW_IMAGES.value and _MODEL_CLASSIFIER_LAYER.value:
+          mnist_lib.plot_images(
+              test_ds,
+              1,
+              5,
+              f"Inference results for {model_descr}",
+              inference_fn=pure_restored_model)
+
+        test_input = np.ones(
+            (mnist_lib.test_batch_size,) + mnist_lib.input_shape,
+            dtype=np.float32)
+        np.testing.assert_allclose(
+            pure_restored_model(tf.convert_to_tensor(test_input)),
+            predict_fn(predict_params, test_input), **tolerances)
+
+  if _SHOW_MODEL.value:
+    def print_model(model_dir: str):
+      cmd = f"saved_model_cli show --all --dir {model_dir}"
+      print(cmd)
+      os.system(cmd)
+
+    print_model(model_dir)
+
+
+def pick_model_class():
+  """Picks one of PureJaxMNIST or FlaxMNIST."""
+  if _MODEL.value == "mnist_pure_jax":
+    return mnist_lib.PureJaxMNIST
+  elif _MODEL.value == "mnist_flax":
+    return mnist_lib.FlaxMNIST
+  else:
+    raise ValueError(f"Unrecognized model: {_MODEL.value}")
+
+
+def model_description() -> str:
+  """A short description of the picked model."""
+  res = pick_model_class().name
+  if not _MODEL_CLASSIFIER_LAYER.value:
+    res += " (features_only)"
+  return res
+
+
+def savedmodel_dir(with_version: bool = True) -> str:
+  """The directory where we save the SavedModel."""
+  model_dir = os.path.join(
+      _MODEL_PATH.value,
+      _MODEL.value + ('' if _MODEL_CLASSIFIER_LAYER.value else '_features')
+  )
+  if with_version:
+    model_dir = os.path.join(model_dir, str(_MODEL_VERSION.value))
+  return model_dir
+
+
+def tf_accelerator_and_tolerances():
+  """Picks the TF accelerator to use and the tolerances for numerical checks."""
+  tf_accelerator = (tf.config.list_logical_devices("TPU") +
+                    tf.config.list_logical_devices("GPU") +
+                    tf.config.list_logical_devices("CPU"))[0]
+  logging.info("Using tf_accelerator = %s", tf_accelerator)
+  if tf_accelerator.device_type == "TPU":
+    tolerances = dict(atol=1e-6, rtol=1e-6)
+  elif tf_accelerator.device_type == "GPU":
+    tolerances = dict(atol=1e-6, rtol=1e-4)
+  elif tf_accelerator.device_type == "CPU":
+    tolerances = dict(atol=1e-5, rtol=1e-5)
+  logging.info("Using tolerances %s", tolerances)
+  return tf_accelerator, tolerances
+
+
+if __name__ == "__main__":
+  app.run(lambda _: train_and_save())

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/examples/saved_model_main_test.py

@@ -0,0 +1,70 @@
+# Copyright 2020 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for mnist_lib, saved_model_lib, saved_model_main."""
+
+import os
+from absl import flags
+from absl.testing import absltest
+from absl.testing import parameterized
+
+from jax._src import config
+from jax._src import test_util as jtu
+
+from jax.experimental.jax2tf.examples import saved_model_main
+from jax.experimental.jax2tf.tests import tf_test_util
+
+config.parse_flags_with_absl()
+FLAGS = flags.FLAGS
+
+
+class SavedModelMainTest(tf_test_util.JaxToTfTestCase):
+
+  def setUp(self):
+    super().setUp()
+    FLAGS.model_path = os.path.join(absltest.get_default_test_tmpdir(),
+                                    "saved_models")
+    FLAGS.num_epochs = 1
+    FLAGS.test_savedmodel = True
+    FLAGS.mock_data = True
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name=f"_{model}_batch={serving_batch_size}",
+          model=model,
+          serving_batch_size=serving_batch_size)
+      for model in ["mnist_pure_jax", "mnist_flax"]
+      for serving_batch_size in [1, -1])
+  def test_train_and_save_full(self,
+                               model="mnist_flax",
+                               serving_batch_size=-1):
+    if (serving_batch_size == -1 and
+        config.jax2tf_default_native_serialization.value and
+        not config.dynamic_shapes.value):
+      self.skipTest("shape polymorphism but --jax_dynamic_shapes is not set.")
+    FLAGS.model = model
+    FLAGS.model_classifier_layer = True
+    FLAGS.serving_batch_size = serving_batch_size
+    saved_model_main.train_and_save()
+
+  @parameterized.named_parameters(
+      dict(testcase_name=f"_{model}", model=model)
+      for model in ["mnist_pure_jax", "mnist_flax"])
+  def test_train_and_save_features(self, model="mnist_flax"):
+    FLAGS.model = model
+    FLAGS.model_classifier_layer = False
+    saved_model_main.train_and_save()
+
+
+if __name__ == "__main__":
+  absltest.main(testLoader=jtu.JaxTestLoader())

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/examples/serving/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2021 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/examples/serving/model_server_request.py

@@ -0,0 +1,128 @@
+# Copyright 2021 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Demonstrates using jax2tf with TensorFlow model server.
+
+See README.md for instructions.
+"""
+import grpc  # type: ignore
+import json
+import logging
+import requests
+
+from absl import app
+from absl import flags
+
+from jax.experimental.jax2tf.examples import mnist_lib
+
+import numpy as np
+import tensorflow as tf
+import tensorflow_datasets as tfds  # type: ignore[import-not-found]
+from tensorflow_serving.apis import predict_pb2  # type: ignore[import-not-found]
+from tensorflow_serving.apis import prediction_service_pb2_grpc
+
+
+_USE_GRPC = flags.DEFINE_boolean(
+    "use_grpc", True,
+    "Use the gRPC API (default), or the HTTP REST API.")
+
+_MODEL_SPEC_NAME = flags.DEFINE_string(
+    "model_spec_name", "",
+    "The name you used to export your model to model server (e.g., mnist_flax).")
+
+_PREDICTION_SERVICE_ADDR = flags.DEFINE_string(
+    "prediction_service_addr",
+    "localhost:8500",
+    "Stubby endpoint for the prediction service. If you serve your model "
+    "locally using TensorFlow model server, then you can use \"localhost:8500\""
+    "for the gRPC server and \"localhost:8501\" for the HTTP REST server.")
+
+_SERVING_BATCH_SIZE = flags.DEFINE_integer(
+    "serving_batch_size",
+    1,
+    "Batch size for the serving request. Must match the "
+    "batch size at which the model was saved. Must divide "
+    "--count_images",
+    lower_bound=1,
+)
+_COUNT_IMAGES = flags.DEFINE_integer(
+    "count_images", 16, "How many images to test.", lower_bound=1
+)
+
+
+def serving_call_mnist(images):
+  """Send an RPC or REST request to the model server.
+
+  Args:
+    images: A numpy.ndarray of shape [B, 28, 28, 1] with the batch of images to
+      perform inference on.
+
+  Returns:
+    A numpy.ndarray of shape [B, 10] with the one-hot inference response.
+  """
+  if _USE_GRPC.value:
+    channel = grpc.insecure_channel(_PREDICTION_SERVICE_ADDR.value)
+    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
+
+    request = predict_pb2.PredictRequest()
+    request.model_spec.name = _MODEL_SPEC_NAME.value
+    request.model_spec.signature_name = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
+    # You can see the name of the input ("inputs") in the SavedModel dump.
+    request.inputs["inputs"].CopyFrom(
+        tf.make_tensor_proto(images, dtype=images.dtype, shape=images.shape))
+    response = stub.Predict(request)
+    # We could also use response.outputs["output_0"], where "output_0" is the
+    # name of the output (which you can see in the SavedModel dump.)
+    # Alternatively, we just get the first output.
+    outputs, = response.outputs.values()
+    return tf.make_ndarray(outputs)
+  else:
+    # Use the HTTP REST api
+    images_json = json.dumps(images.tolist())
+    # You can see the name of the input ("inputs") in the SavedModel dump.
+    data = f'{{"inputs": {images_json}}}'
+    predict_url = f"http://{_PREDICTION_SERVICE_ADDR.value}/v1/models/{_MODEL_SPEC_NAME.value}:predict"
+    response = requests.post(predict_url, data=data)
+    if response.status_code != 200:
+      msg = (f"Received error response {response.status_code} from model "
+             f"server: {response.text}")
+      raise ValueError(msg)
+    outputs = response.json()["outputs"]
+    return np.array(outputs)
+
+
+def main(_):
+  if _COUNT_IMAGES.value % _SERVING_BATCH_SIZE.value != 0:
+    raise ValueError(f"The count_images ({_COUNT_IMAGES.value}) must be a "
+                     "multiple of "
+                     f"serving_batch_size ({_SERVING_BATCH_SIZE.value})")
+  test_ds = mnist_lib.load_mnist(tfds.Split.TEST,
+                                 batch_size=_SERVING_BATCH_SIZE.value)
+  images_and_labels = tfds.as_numpy(test_ds.take(
+      _COUNT_IMAGES.value // _SERVING_BATCH_SIZE.value))
+
+  accurate_count = 0
+  for batch_idx, (images, labels) in enumerate(images_and_labels):
+    predictions_one_hot = serving_call_mnist(images)
+    predictions_digit = np.argmax(predictions_one_hot, axis=1)
+    labels_digit = np.argmax(labels, axis=1)
+    accurate_count += np.sum(labels_digit == predictions_digit)
+    running_accuracy = (
+        100. * accurate_count / (1 + batch_idx) / _SERVING_BATCH_SIZE.value)
+    logging.info(
+        " predicted digits = %s labels %s. Running accuracy %.3f%%",
+        predictions_digit, labels_digit, running_accuracy)
+
+
+if __name__ == "__main__":
+  app.run(main)

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/impl_no_xla.py

@@ -0,0 +1,1287 @@
+# Copyright 2020 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Workarounds for jax2tf transforms when XLA is not linked in."""
+
+from __future__ import annotations
+
+import builtins
+from collections.abc import Sequence
+import dataclasses
+from functools import partial, wraps
+import math
+import string
+from typing import Any, Callable, Optional
+
+from jax._src import core
+from jax import lax
+from jax._src.lax import slicing as lax_slicing
+from jax._src import dtypes
+from jax._src import util
+
+from jax.experimental.jax2tf import jax2tf
+
+import numpy as np
+import tensorflow as tf
+
+
+# Implementation rules for primitives when XLA is not linked in. These
+# implementations are workarounds, making use of TF ops that do work when XLA is
+# not linked in. They are only used when the argument `enable_xla=False` when
+# calling jax2tf.convert().
+tf_impl_no_xla: dict[core.Primitive, Callable[..., Any]] = {}
+
+
+TfVal = Any
+DType = Any
+PrecisionType = Any
+
+
+def _error(primitive_name: str, suffix_msg: str = "") -> Exception:
+  msg = f"Call to {primitive_name} cannot be converted with enable_xla=False."
+  if suffix_msg:
+    msg += (f" {suffix_msg} - See source code for the precise conditions under "
+             "which it can be converted without XLA.")
+  return NotImplementedError(msg)
+
+_conv_error = lambda msg: _error("conv_general_dilated", msg)
+_reduce_error = lambda msg: _error("reduce_window", msg)
+_scatter_error = lambda msg: _error("scatter_(update/add/multiply/min/max)", msg
+                                   )
+
+def _unimplemented(name):
+
+  def op(*arg, **kwargs):
+    raise _error(name)
+
+  return op
+
+
+# TODO(marcvanzee): Remove this function and use `tf.math.invert_permutation`
+# once it is implemented by TFjs:
+# https://github.com/tensorflow/tfjs/issues/6395.
+def _invert_permutation(perm):
+  return tuple(perm.index(i) for i in range(len(perm)))
+
+
+def _transpose_with_shape(x: TfVal, x_shape: core.Shape, permutation) -> tuple[TfVal, core.Shape]:
+  """Computes transposition of x and its shape.
+
+  x_shape matches x.shape in the known dimensions, and it has dimension
+  polynomials elsewhere, while x.shape has None.
+  """
+  return tf.transpose(x, perm=permutation), tuple(x_shape[i] for i in permutation)
+
+
+def _transpose_for_tf_conv(lhs, lhs_shape: core.Shape,
+                           rhs, rhs_shape: core.Shape, dimension_numbers):
+  """Transposes lhs and rhs to respectively NHWC and HWIO so they can be passed to TF functions.
+
+  The shapes passed in and returned may contain polynomials, and thus may
+  be different than lhs.shape and rhs.shape.
+  """
+  # TODO(marcvanzee): Add tests for this ops for shape polymorphism.
+  lhs_perm, rhs_perm, _ = dimension_numbers
+
+  # TODO(marcvanzee): Consider merging transposes if we want to optimize.
+  # For `lhs_perm` / `output_perm`, perm (0, 1, 2, 3) corresponds to "NCHW".
+  lhs, lhs_shape = _transpose_with_shape(lhs, lhs_shape, lhs_perm)  # lhs --> "NCHW"
+  if len(lhs_perm) == 3:
+    # For 1D convolution, we add a trivial "W" dimension, so that 2D Convolution
+    # logic can be applied downstream.
+    lhs = lhs[:, :, :, np.newaxis]
+    lhs_shape = tuple(lhs_shape) + (1,)
+  # However, the TF ops only support "NHWC" on CPU, so we transpose again.
+  lhs, lhs_shape = _transpose_with_shape(lhs, lhs_shape, (0, 2, 3, 1))  # "NCHW" --> "NHWC"
+
+  # For `rhs_perm`, perm (0, 1, 2, 3) corresponds to "OIHW".
+  rhs, rhs_shape = _transpose_with_shape(rhs, rhs_shape, rhs_perm)  # rhs --> "OIHW"
+  # Handle conv1d case.
+  if len(rhs_perm) == 3:
+    rhs = rhs[:, :, :, np.newaxis]
+    rhs_shape = tuple(rhs_shape) + (1,)
+  # For the tf ops, rhs is expected to be "OIHW".
+  rhs, rhs_shape = _transpose_with_shape(rhs, rhs_shape, (2, 3, 1, 0))  # "OIHW" --> "HWIO"
+  jax2tf._assert_matching_abstract_shape(lhs, lhs_shape)
+  jax2tf._assert_matching_abstract_shape(rhs, rhs_shape)
+  return lhs, lhs_shape, rhs, rhs_shape
+
+
+def pads_to_padtype(in_shape, window_shape, window_strides, padding) -> str:
+  for pad_str in ["VALID", "SAME"]:
+    pads = lax.padtype_to_pads(in_shape, window_shape, window_strides, pad_str)
+    if list(pads) == list(padding):
+      return pad_str
+  return "EXPLICIT"
+
+
+def _pad_spatial_dims(x, x_shape, padding):
+  """Pads `x` using `padding`, which specifies padding for the spatial dimensions."""
+  padding = tuple(padding)
+  if len(padding) == len(x_shape) - 2:
+    # If necessary, add empty padding for batch and feature dimensions.
+    no_pad = ((0, 0),)
+    padding = no_pad + padding + no_pad
+  x = tf.pad(x, padding)
+  assert len(x.shape) == len(padding)
+  x_shape = tuple(p0 + xs + p1 for xs, (p0, p1) in zip(x_shape, padding))
+  jax2tf._assert_matching_abstract_shape(x, x_shape)
+  return x, x_shape
+
+def _check_pad_spatial_dims(x, x_shape, padding):
+  """Pads `x` using `padding`, which specifies padding for the spatial dimensions."""
+  padding = tuple(padding)
+  if len(padding) == len(x_shape) - 2:
+    # If necessary, add empty padding for batch and feature dimensions.
+    no_pad = ((0, 0),)
+    padding = no_pad + padding + no_pad
+  assert len(x.shape) == len(padding)
+  x_shape = tuple(p0 + xs + p1 for xs, (p0, p1) in zip(x_shape, padding))
+  return x, x_shape, padding
+
+def _conv_transpose_pads_to_padtype(kernel_sdims, lhs_dilation, padding):
+  """Finds the padding type for a transpose convolution."""
+  # This is simply checking agreement with lax._conv_transpose_padding.
+  is_valid = True
+  is_same = True
+  if not len(kernel_sdims) == len(lhs_dilation) == len(padding):
+    raise ValueError(f'Found different lengths for '
+                     f'kernel_sdims ({kernel_sdims}), '
+                     f'lhs_dilation ({lhs_dilation}), '
+                     f'and padding ({padding}).')
+  for k, s, (begin, end) in zip(kernel_sdims, lhs_dilation, padding):
+    # Check for VALID padding.
+    pad_len_valid = k + s - 2 + builtins.max(k - s, 0)
+    pad_a = k - 1
+    pad_b = pad_len_valid - pad_a
+    if begin != pad_a or end != pad_b:
+      is_valid = False
+
+    # Check for SAME padding.
+    pad_len_same = k + s - 2
+    if s > k - 1:
+      pad_a = k - 1
+    else:
+      pad_a = int(np.ceil(pad_len_same / 2))
+    pad_b = pad_len_same - pad_a
+    if begin != pad_a or end != pad_b:
+      is_same = False
+
+  if is_valid:
+    return 'VALID'
+  elif is_same:
+    return 'SAME'
+  raise ValueError('Transpose convolution padding mode must be '
+                   '`SAME` or `VALID`.')
+
+def _validate_spatial_dimensions(lhs: TfVal, lhs_shape: core.Shape,
+                                 rhs: TfVal, rhs_shape: core.Shape):
+  """Check spatial dimension support."""
+  jax2tf._assert_matching_abstract_shape(lhs, lhs_shape)
+  jax2tf._assert_matching_abstract_shape(rhs, rhs_shape)
+
+  nr_spatial_dimensions = len(lhs_shape) - 2
+  # Currently we only support 1D+2D convolutions because it keeps the code
+  # relatively simple and covers most cases.
+  if nr_spatial_dimensions > 2:
+    raise _conv_error(
+        "We only support 1D or 2D convolutions, but found "
+        f"{nr_spatial_dimensions}.")
+
+
+def _normalize_padding_and_dilations(
+      padding, lhs_dilation, rhs_dilation, is_conv1d):
+  if is_conv1d:
+    lhs_dilation = list(lhs_dilation) + [1]
+    rhs_dilation = list(rhs_dilation) + [1]
+    # Empty padding in the dummy dimension.
+    # Note that when kernel_size=stride=1, padding of (0, 0) is both 'VALID' and
+    # 'SAME'. So the inferred padding type will still register according to the
+    # first dimension padding.
+    padding = list(padding) + [(0, 0)]
+  return padding, lhs_dilation, rhs_dilation
+
+
+def _normalize_window_strides(window_strides):
+  """Ensure window_strides has length 4."""
+  # Some TF ops require len(window_strides) == 4 while others do not. We simply
+  # ensure it always has len(4).
+  if len(window_strides) == 1:
+    # This is the Conv1D case. We add a dummy dimension to allow using 2D ops,
+    # and use stride=1 on the dummy dimension.
+    window_strides = list(window_strides) + [1]
+  if len(window_strides) == 2:
+    window_strides = [1] + list(window_strides) + [1]
+  return window_strides
+
+
+def _validate_conv_features(
+      is_transpose, is_atrous, is_depthwise, feature_group_count,
+      batch_group_count, preferred_element_type, lhs_dtype):
+  if feature_group_count > 1 and not is_depthwise:
+    raise _conv_error("Grouped convolutions are unsupported")
+  if (is_depthwise and is_atrous) and not is_transpose:
+    # We allow dilated depthwise convolutions.
+    pass
+  elif [is_depthwise, is_atrous, is_transpose].count(True) > 1:
+    raise _conv_error(
+        f"Can only do one of depthwise ({is_depthwise}), atrous ({is_atrous}) "
+        f"and transposed convolutions ({is_transpose})")
+
+  # We can implement batch grouping when there is a need for it.
+  if batch_group_count != 1:
+    raise _conv_error("Unimplemented support for batch_group_count != 1 "
+                f"(found {batch_group_count})")
+
+  if (preferred_element_type is not None and
+      preferred_element_type != lhs_dtype):
+    raise _conv_error("Unimplemented support for preferred_element_type")
+
+
+def _conv_general_dilated(
+    lhs, rhs, *, window_strides, padding, lhs_dilation, rhs_dilation,
+    dimension_numbers: lax.ConvDimensionNumbers, feature_group_count: int,
+    batch_group_count: int,
+    precision: tuple[PrecisionType, PrecisionType] | None,
+    preferred_element_type: DType | None,
+    _in_avals: Sequence[core.ShapedArray], _out_aval: core.ShapedArray):
+  """Implementation of lax.conv_general_dilated_p using XlaConv."""
+  # In presence of shape polymorphism, lhs.shape and rhs.shape may contain
+  # None. The actual dimension polynomial shapes are in _in_avals.
+  del precision  # Unused arguments.
+  lhs_shape, rhs_shape = _in_avals[0].shape, _in_avals[1].shape
+  out_shape = _out_aval.shape
+  _validate_spatial_dimensions(lhs, lhs_shape, rhs, rhs_shape)
+  is_conv1d = len(lhs_shape) - 2 == 1
+
+  tf_window_strides = _normalize_window_strides(window_strides)
+  padding, lhs_dilation, rhs_dilation = _normalize_padding_and_dilations(
+      padding, lhs_dilation, rhs_dilation, is_conv1d)
+
+  lhs, lhs_shape, rhs, rhs_shape = _transpose_for_tf_conv(lhs, lhs_shape,
+                                                          rhs, rhs_shape,
+                                                          dimension_numbers)
+  in_channels = lhs_shape[-1]
+  *rhs_spatial_shapes, _, rhs_out_channel = rhs_shape
+
+  is_transpose = any(d != 1 for d in lhs_dilation)
+  is_atrous = any(d != 1 for d in rhs_dilation)
+  is_depthwise = in_channels == feature_group_count and feature_group_count > 1
+  _validate_conv_features(is_transpose, is_atrous, is_depthwise,
+                          feature_group_count, batch_group_count,
+                          preferred_element_type, lhs.dtype.as_numpy_dtype)
+
+  rhs_dilated_shape = [
+      (k - 1) * r + 1 for k, r in zip(rhs_spatial_shapes, rhs_dilation)
+  ]
+  output_perm = dimension_numbers[2]
+
+  if is_transpose:
+    padding_type = _conv_transpose_pads_to_padtype(
+        rhs_spatial_shapes, lhs_dilation, padding)
+  else:
+    padding_type = pads_to_padtype(
+      lhs_shape[1:3], rhs_dilated_shape, window_strides, padding)
+    # We only manually pad if we aren't using a transposed convolutions.
+    if padding_type == "EXPLICIT":
+      lhs, lhs_shape, padding = _check_pad_spatial_dims(lhs, lhs_shape, padding)
+      padding_type = padding
+
+  if padding_type != "SAME" and any(l < r for l, r in zip(lhs_shape[1:3], rhs_dilated_shape)):
+    # If the input shape is smaller than the filter shape in a spatial dimension,
+    # lax returns only zeros while tf.conv2d returns an error.
+    # We thus return zeros to make sure the behavior is consistent.
+    return tf.broadcast_to(tf.constant(0, dtype=tf.float32),
+                           jax2tf._eval_shape(out_shape))
+
+  if is_depthwise:
+    # Reshape filter from
+    # [filter_height, filter_width, 1, in_channels * channel_multiplier] to
+    # [filter_height, filter_width, in_channels, channel_multiplier].
+    new_rhs_shape = tuple(rhs_spatial_shapes) + (in_channels,
+                                                 rhs_out_channel // in_channels)
+    output = tf.nn.depthwise_conv2d(
+        input=lhs,
+        filter=tf.reshape(rhs, jax2tf._eval_shape(new_rhs_shape)),
+        strides=tf_window_strides,
+        padding=padding_type,
+        dilations=rhs_dilation)
+
+  elif is_transpose:
+    # tf.nn.conv2d_transpose requires a transposed filter.
+    rhs_t = tf.reverse(rhs, [0, 1])
+    rhs_t = tf.transpose(rhs_t, (0, 1, 3, 2))
+
+    # We should transpose `out_shape` to "NHWC", which is what TF expects.
+    # First transpose to "NCHW".
+    if is_conv1d:
+      tf_out_shape = tuple(out_shape[i] for i in output_perm) + (1,)
+    else:
+      tf_out_shape = tuple(out_shape[i] for i in output_perm)
+    # Then transpose "NCHW" to "NHWC".
+    tf_out_shape = tuple(tf_out_shape[i] for i in (0, 2, 3, 1))
+    output = tf.nn.conv2d_transpose(
+        input=lhs,
+        filters=rhs_t,
+        output_shape=jax2tf._eval_shape(tf_out_shape),
+        strides=lhs_dilation,
+        padding=padding_type)
+
+  else:
+    output = tf.nn.conv2d(
+        input=lhs,
+        filters=rhs,
+        strides=tf_window_strides,
+        padding=padding_type,
+        dilations=rhs_dilation)
+
+  # TF outputs in format "NHWC", so convert to "NCHW", which is lax's default
+  # format.
+  output = tf.transpose(output, (0, 3, 1, 2))  # "NHWC" --> "NCHW"
+  if is_conv1d:
+    output = output[:, :, :, 0]
+  # To determine the right permutation, we compute the inverse permutation of
+  # `output_perm`, so that when `output_perm` is applied to `output`, we obtain
+  # the outpt in NCHW format.
+  inverse_perm = _invert_permutation(output_perm)
+  output = tf.transpose(output, inverse_perm)  # "NCHW" -> desired output shape.
+  return output
+
+
+tf_impl_no_xla[lax.conv_general_dilated_p] = _conv_general_dilated
+
+
+def _dot_general(lhs, rhs, *, dimension_numbers,
+                 precision: tuple[PrecisionType, PrecisionType] | None,
+                 preferred_element_type: DType | None,
+                 _in_avals: Sequence[core.ShapedArray],
+                 _out_aval: core.ShapedArray):
+  """Implementation of lax.dot_general_p in terms of tf.linalg.einsum."""
+  # Unused arguments.
+  del precision
+  del preferred_element_type
+
+  lhs, rhs, convert_result = jax2tf._dot_general_convert_to_common_dtype(
+    lhs, _in_avals[0], rhs, _in_avals[1], _out_aval)
+
+  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = dimension_numbers
+  lhs_ndim, rhs_ndim = len(lhs.shape), len(rhs.shape)
+
+  # This condition ensures that:
+  # 1) the batch dimensions are ordered in the same way in lhs and rhs (this is
+  #    not strictly necessary, but we would have to reshape the array if that
+  #    were not the case;
+  # 2) lhs and rhs have the same number of dimensions +/- 1
+  # 3) the number of non-batch dimensions in both tensors is either 1 or 2
+  # 4) the contracting dimensions are consistent with those of a classic
+  #    matrix/matrix, vector/matrix or matrix/vector multiplication.
+  if (lhs_batch == rhs_batch == tuple(range(len(lhs_batch))) and
+      lhs_ndim - rhs_ndim in [-1, 0, 1] and
+      1 <= lhs_ndim - len(lhs_batch) <= 2 and
+      1 <= rhs_ndim - len(rhs_batch) <= 2 and
+      lhs_contracting == (len(lhs.shape) - 1,) and
+      rhs_contracting == (len(lhs_batch),)):
+    # All the inputs to tf.linalg.matmul must have 2 inner dimensions,
+    # after their batch dimensions, so we need to expand the dimensions
+    # appropriately. We can get to this branch with three combinations of
+    # inner shapes:
+    # - lhs.inner_shape == [a, b], rhs.inner_shape == [b, c]
+    #   - in this case, the resulting inner shape is [a, c];
+    # - lhs.inner_shape == [b]   , rhs.inner_shape == [b, c]
+    #   - in this case, we need to expand lhs to [1, b], and the resulting
+    #     shape is [c]. We need to squeeze the result of tf.linalg.matmul
+    #     as it will have shape [1, c];
+    # - lhs.shape == [batch] + [a, b], rhs.shape == [batch] + [b]
+    #   - in this case, we need to expand rhs to [b, 1], and the resulting
+    #     shape is [a]. We need to squeeze the result of tf.linalg.matmul
+    #     as it will have shape [a, 1];
+    # - lhs.shape == [batch] + [b]   , rhs.shape == [batch] + [b]
+    #   - in this case, we need to expand lhs to [1, b] and rhs to [b, 1],
+    #     and the resulting shape is (). We need to squeeze the result of
+    #     tf.linalg.matmul as it will have shape [1, 1].
+    squeeze_idxs = []
+    if lhs_ndim - len(lhs_batch) == 1:
+      lhs = tf.expand_dims(lhs, lhs_ndim - 1)
+      squeeze_idxs.append(len(lhs.shape) - 2)
+    if rhs_ndim - len(rhs_batch) == 1:
+      rhs = tf.expand_dims(rhs, rhs_ndim)
+      squeeze_idxs.append(len(rhs.shape) - 1)
+    result = tf.linalg.matmul(lhs, rhs)
+    if len(squeeze_idxs) != 0:
+      assert all(result.shape[i] == 1 for i in squeeze_idxs)
+      result = tf.squeeze(result, squeeze_idxs)
+    return convert_result(result)
+
+  new_id = iter(string.ascii_letters)
+  lhs_axis_ids = [next(new_id) for _ in lhs.shape]
+  rhs_axis_ids = [next(new_id) for _ in rhs.shape]
+  lhs_out_axis_ids = lhs_axis_ids[:]
+  rhs_out_axis_ids = rhs_axis_ids[:]
+
+  for lhs_axis, rhs_axis in zip(lhs_contracting, rhs_contracting):
+    shared_id = next(new_id)
+    lhs_axis_ids[lhs_axis] = shared_id
+    rhs_axis_ids[rhs_axis] = shared_id
+    lhs_out_axis_ids[lhs_axis] = None  # type: ignore[call-overload]
+    rhs_out_axis_ids[rhs_axis] = None  # type: ignore[call-overload]
+
+  batch_ids = []
+  for lhs_axis, rhs_axis in zip(lhs_batch, rhs_batch):
+    shared_id = next(new_id)
+    lhs_axis_ids[lhs_axis] = shared_id
+    rhs_axis_ids[rhs_axis] = shared_id
+    lhs_out_axis_ids[lhs_axis] = None  # type: ignore[call-overload]
+    rhs_out_axis_ids[rhs_axis] = None  # type: ignore[call-overload]
+    batch_ids.append(shared_id)
+
+  not_none = lambda x: x is not None
+  out_axis_ids = list(
+      filter(not_none, batch_ids + lhs_out_axis_ids + rhs_out_axis_ids))
+  assert lhs.dtype == rhs.dtype
+  spec = "{},{}->{}".format("".join(lhs_axis_ids), "".join(rhs_axis_ids),
+                            "".join(out_axis_ids))
+  return convert_result(tf.linalg.einsum(spec, lhs, rhs))
+
+
+tf_impl_no_xla[lax.dot_general_p] = _dot_general
+
+
+def _interior_padding(operand, padding_value, padding_config, operand_shape):
+  # Used only when enable_xla=False
+  # Applies only the interior padding from the padding_config.
+  # We do this somewhat inefficiently, as a scatter.
+  # For each dimension we compute the indices_by_dim as [0, f, 2f, 3f, ...] where
+  # f is the dilation factor for the dimension, i.e., 1 + interior_padding.
+  # Then we compute the cartesian production of the indices (using broadcast
+  # and concat).
+
+  # We could make this code more complex and do all the padding at once, but
+  # we prefer to keep it simple.
+  indices_by_dim = []
+  indices_shape = operand_shape + (1,)
+  output_shape = []  # considering only interior padding
+  for d, (dsz, (_, _, i)) in enumerate(zip(operand_shape, padding_config)):
+    dilation_factor = i + 1
+    output_shape.append(dsz * dilation_factor - i)
+    indices = tf.range(dsz) * dilation_factor
+    expansion = [None] * (1 + len(operand_shape))
+    expansion[d] = slice(None, None, None)
+    indices_by_dim.append(tf.broadcast_to(indices[expansion], indices_shape))
+
+  indices_cartesian = tf.concat(indices_by_dim, axis=len(operand_shape))
+  scattered = tf.scatter_nd(indices_cartesian, operand, output_shape)
+  # What elements from the output array we use from
+  mask = tf.scatter_nd(indices_cartesian, tf.ones_like(operand, dtype=np.bool_),
+                       output_shape)
+  return tf.where(mask, scattered, padding_value)
+
+
+def _pad(operand, padding_value, *, padding_config,
+         _in_avals: Sequence[core.ShapedArray], _out_aval: core.ShapedArray):
+  # Do only the interior padding first. This is rarely needed.
+  if any(i != 0 for _, _, i in padding_config):
+    operand = _interior_padding(operand, padding_value, padding_config,
+                                jax2tf._eval_shape(_in_avals[0].shape))
+
+  # Now do the non-negative edge padding. This is the common case, use tf.pad.
+  non_negative_padding = [((lo if lo >= 0 else 0), (hi if hi >= 0 else 0))
+                          for lo, hi, _ in padding_config]
+  operand = tf.pad(
+      operand,
+      non_negative_padding,
+      mode="CONSTANT",
+      constant_values=padding_value)
+  # Now the negative edge padding (this is also rare)
+  if any(lo < 0 or hi < 0 for lo, hi, _ in padding_config):
+    output_shape = jax2tf._eval_shape(_out_aval.shape)
+    begins = [(-lo if lo < 0 else 0) for lo, _, _ in padding_config]
+    operand = tf.slice(operand, begins, output_shape)
+
+  return operand
+
+
+tf_impl_no_xla[lax.pad_p] = _pad
+
+
+def _argminmax(is_min: bool, operand: TfVal, axes: Sequence[int],
+               index_dtype: DType, _in_avals: Sequence[core.ShapedArray],
+               _out_aval: core.ShapedArray):
+  # The following is known to diverge from JAX behavior for NaN.
+  axis, = axes
+  output_type = tf.int32
+  if dtypes.iinfo(index_dtype).bits > 32:
+    output_type = tf.int64
+  # TODO(phawkins): handle axes larger than 2^31.
+  fn = tf.math.argmin if is_min else tf.math.argmax
+  result = fn(operand, axis=axis, output_type=output_type)
+  return tf.cast(result, jax2tf._to_tf_dtype(index_dtype))
+
+
+tf_impl_no_xla[lax.argmin_p] = partial(_argminmax, True)
+tf_impl_no_xla[lax.argmax_p] = partial(_argminmax, False)
+
+
+def _validate_reduce_window_inputs(operand_shape, computation_name, dtype,
+                                   window_dimensions, window_strides,
+                                   base_dilation, window_dilation):
+  if computation_name not in ["min", "max", "add"]:
+    raise _reduce_error("Reduction function should be either min, max, or add.")
+  if computation_name in ["min", "max"] and dtype in [
+      tf.bool, tf.uint32, tf.uint64, tf.complex64, tf.complex128
+  ]:
+    raise _reduce_error("Min/max pool does not support operands of type "
+                        f"{dtype}")
+  if computation_name == "min" and dtype in [tf.uint8, tf.uint16]:
+    # TODO(marcvanzee): We currently implement min pooling by negating the
+    # input, but this doesn't work for uint. We could work around it using
+    # tf.math.reduce_min.
+    raise _reduce_error(f"Min pool does not support operands of type {dtype}")
+  if computation_name == "add" and dtype not in [
+      tf.bfloat16,
+      tf.float16,
+      tf.float32,
+      tf.float64,
+      tf.int16,
+      tf.int32,
+  ]:
+    raise _reduce_error("Add pooling does not support operands of type "
+                        f"{dtype}")
+
+  if (len(operand_shape) != len(window_dimensions) != len(window_strides) !=
+      len(window_dilation)):
+    raise _reduce_error("Input shapes, window dimensions, window stride "
+                        "dimensions, and window dilation dimensions should "
+                        "match.")
+
+  has_only_spatial_dims = True
+  if len(operand_shape) > 4:
+    raise _reduce_error("Only 1D or 2D input are supported.")
+  if len(operand_shape) > 2:
+    # operand_shape = (batch, spatial_dims, ..., channel).
+    has_only_spatial_dims = False
+
+    for name, value in [("window_dimensions", window_dimensions),
+                        ("window_strides", window_strides),
+                        ("window_dilation", window_dilation)]:
+      if value[0] != value[-1] != 1:
+        raise _reduce_error("Only 1D or 2D input are supported, expected "
+                            f"{name}=(1, spatial_dims, ..., 1), but got "
+                            f"{value}")
+
+  if list(base_dilation) != [1] * len(operand_shape):
+    # TODO(marcvanzee): Add support for base dilations. We can do this using
+    # a scatter on operand.
+    raise _reduce_error("Unimplemented support for base dilation.")
+
+  return has_only_spatial_dims
+
+
+def _padding_reduce_window(operand, operand_shape, computation_name,
+                           window_dimensions, window_strides, padding):
+  padding_type = pads_to_padtype(operand_shape, window_dimensions,
+                                 window_strides, padding)
+
+  # https://github.com/google/jax/issues/11874.
+  needs_manual_padding = (
+      padding_type == "SAME" and computation_name == "add" and
+      window_dimensions != [1] * len(operand_shape))
+
+  if needs_manual_padding or padding_type == "EXPLICIT":
+    operand, operand_shape = _pad_spatial_dims(operand, operand_shape, padding)
+    padding_type = "VALID"
+
+  return operand, operand_shape, padding_type
+
+
+def _reshape_reduce_window(operand, operand_shape, window_dimensions,
+                           window_strides, window_dilation, *,
+                           has_only_spatial_dims):
+  # Reshape inputs so they are accepted by tf.nn.pool, which expects batch and
+  # channel dimensions for operand but not for any of the other inputs.
+  if has_only_spatial_dims:  # len(operand_shape) <= 2
+    # Call eval_shape on a shape that may contain polynomials, otherwise TF does
+    # not know what to do with polynomials in the shape.
+    operand_shape = jax2tf._eval_shape(operand_shape)
+    # Add batch and channel dimensions to operand.
+    operand = tf.reshape(operand, (1,) + operand_shape + (1,))
+  else:
+    # This branch assumes operand.shape = (batch, spatial_dims, ..., channel),
+    # and dimensions, strides, dilation are all (1, spatial_values, ..., 1).
+    # Input validation for this is done in _validate_reduce_window_inputs.
+    window_dimensions = window_dimensions[1:-1]
+    window_strides = window_strides[1:-1]
+    window_dilation = window_dilation[1:-1]
+
+  return operand, window_dimensions, window_strides, window_dilation
+
+
+def _reduce_monoid(operand, window_dimensions, window_strides, padding,
+                   base_dilation, window_dilation, computation_name,
+                   _in_avals: Sequence[core.ShapedArray],
+                   _out_aval: core.ShapedArray):
+  dtype = operand.dtype
+  # In presence of shape polymorphism, operand.shape may contain None. The
+  # actual dimension polynomial shapes are in _in_avals.
+  operand_shape = _in_avals[0].shape
+
+  # TODO(marcvanzee): Put reduce_window arguments into dataclass, similar to
+  # Gather, to simplify function calls.
+  has_only_spatial_dims = _validate_reduce_window_inputs(
+      operand_shape, computation_name, dtype, window_dimensions, window_strides,
+      base_dilation, window_dilation)
+
+  operand, operand_shape, padding_type = _padding_reduce_window(
+      operand, operand_shape, computation_name, window_dimensions,
+      window_strides, padding)
+
+  operand, window_dimensions, window_strides, dilations = _reshape_reduce_window(
+      operand,
+      operand_shape,
+      window_dimensions,
+      window_strides,
+      window_dilation,
+      has_only_spatial_dims=has_only_spatial_dims)
+
+  def tf_pool(inputs, pooling_type):
+    if any(not core.is_constant_shape(s) for s in
+           (window_dimensions, window_strides, dilations)):
+      raise NotImplementedError(
+          f"TODO: use tf.nn.pool with dynamic shapes¨{window_dimensions=} "
+          f" {window_strides=} {dilations=}")
+    # tf.nn.pool() currently does not suport tf.int32 and so we cast back and
+    # forth in order to be able to convert.
+    if (inputs.dtype in [tf.int16, tf.int32]) and computation_name == "add":
+      original_dtype = inputs.dtype
+      inputs = tf.cast(inputs, dtype=tf.float32)
+    else:
+      original_dtype = None
+    result = tf.nn.pool(
+        inputs,
+        window_shape=window_dimensions,
+        pooling_type=pooling_type,
+        padding=padding_type,
+        strides=window_strides,
+        dilations=dilations)
+    if original_dtype:
+      result = tf.cast(result, dtype=original_dtype)
+
+    if has_only_spatial_dims:
+      # If the input only had spatial dimensions we need to contract the batch
+      # and channel dimensions before returning the output.
+      result = tf.squeeze(result, [0, -1])
+
+    jax2tf._assert_matching_abstract_shape(result, _out_aval.shape)
+    return result
+
+  negate = lambda x: tf.multiply(x, tf.constant(-1, dtype))
+  if computation_name == "max":
+    return tf_pool(operand, "MAX")
+  elif computation_name == "min":
+    return negate(tf_pool(negate(operand), "MAX"))
+  elif computation_name == "add":
+    # TODO(marcvanzee): This may give very large deviations on TPU when using
+    # floats as inputs. Alternatively, we could implement this using a
+    # convolution with an all-1's kernel.
+    return tf.multiply(tf_pool(operand, "AVG"), math.prod(window_dimensions))
+
+
+def _reduce_window(*args, jaxpr, consts, window_dimensions,
+                   window_strides, padding, base_dilation, window_dilation,
+                   _in_avals: Sequence[core.ShapedArray],
+                   _out_aval: tuple[core.ShapedArray, ...]
+                   ) -> tuple[TfVal, ...]:
+  assert len(consts) == 0, "Reduction computation cannot have constants"
+  operands, init_values = util.split_list(args, [len(args) // 2])
+
+  if len(operands) != 1 or len(init_values) != 1:
+    raise _reduce_error("jax2tf does not support variadic reduce_window")
+
+  operand, init_value = operands[0], init_values[0]
+  # Infer operation type from jaxpr.
+  if (len(jaxpr.eqns) != 1 or
+      len(jaxpr.eqns[0].invars) != 2 or
+      len(jaxpr.eqns[0].outvars) != 1 or
+      jaxpr.eqns[0].primitive.name not in ["min", "max", "add"]):
+    raise _reduce_error("Reduction function should be either min, max, or add.")
+
+  computation_name = jaxpr.eqns[0].primitive.name
+  result = _reduce_monoid(operand,
+                          window_dimensions=window_dimensions,
+                          window_strides=window_strides,
+                          padding=padding,
+                          base_dilation=base_dilation,
+                          window_dilation=window_dilation,
+                          computation_name=computation_name,
+                          _in_avals=(_in_avals[0],),  # Don't pass init_value.
+                          _out_aval=_out_aval[0])     # Returns single value.
+
+  reduce_fn = {
+      "min": tf.minimum,
+      "max": tf.maximum,
+      "add": tf.add,
+  }[computation_name]
+  result = reduce_fn(result, init_value)
+
+  # The output is expected to be wrapped in a tuple, and since we don't use
+  # variadic reductions, this tuple always contains a single element.
+  return (result,)
+
+
+tf_impl_no_xla[lax.reduce_window_min_p] = (
+    partial(_reduce_monoid, computation_name="min"))
+tf_impl_no_xla[lax.reduce_window_max_p] = (
+    partial(_reduce_monoid, computation_name="max"))
+tf_impl_no_xla[lax.reduce_window_sum_p] = (
+    partial(_reduce_monoid, computation_name="add"))
+
+tf_impl_no_xla[lax.reduce_window_p] = _reduce_window
+
+tf_impl_no_xla[lax.reduce_p] = _unimplemented("reduce")
+
+tf_impl_no_xla[lax.select_and_scatter_add_p] = _unimplemented(
+    "select_and_scatter_add")
+
+tf_impl_no_xla[lax.rng_bit_generator_p] = _unimplemented("rng_bit_generator")
+
+
+def _clip(max_indices: Sequence[TfVal], start_indices: Sequence[TfVal],
+          slice_sizes: Sequence[TfVal]):
+  """Simulates XLA clipping behavior with TF ops.
+
+  Various TF ops have different clipping behavior than XLA:
+  * If `start_indices` is out-of-bounds, then TF fails but XLA clips the indices
+  to
+    [0, max_len].
+  * If `start_indices + slice_size` is out-of-bounds, then TF fails, but XLA
+  adjust
+    `start_indices` so that a full slice is returned.
+  This function clips the start indices correctly.
+  """
+  # We cast both arguments to `tf.clip_by_value` to int32. Otherwise, this
+  # function may return uint32 which is not always compatible with TF ops, so
+  # this may result in type errors.
+  max_start = tf.cast(tf.subtract(max_indices, slice_sizes), dtype=tf.int32)
+  return tf.clip_by_value(tf.cast(start_indices, dtype=tf.int32), 0, max_start)
+
+
+@dataclasses.dataclass
+class GatherArgs:
+  operand: TfVal
+  start_indices: TfVal
+  dnums: lax.GatherDimensionNumbers
+  slice_sizes: TfVal
+  op_shape: core.Shape
+  start_indices_shape: core.Shape
+  out_aval: core.ShapedArray
+
+  def __post_init__(self):
+    assert len(self.op_shape) == len(self.slice_sizes)
+
+  def __repr__(self):
+    return (f"operand shape={self.op_shape}, "
+            f"start_indices={self.start_indices}, "
+            f"dimension_numbes={self.dnums}, "
+            f"slice_sizes={self.slice_sizes}")
+  @property
+  def batch_dims(self):
+    return tuple(x for x in range(len(self.out_aval.shape))
+                if x not in self.dnums.offset_dims)
+
+def gather_precondition(precondition_fn: Callable[[GatherArgs], None]):
+  """Decorator for specifying a precondition function.
+
+  This decorator should be put on a function with argument `arg` of type
+  `GatherArgs`. It will first call `precondition_fn` with `arg` (which may throw
+  an exception), and then call the function it is decorating with `arg` as well.
+  """
+
+  def decorator(gather_fn: Callable[[GatherArgs], Any]):
+
+    @wraps(gather_fn)
+    def wrapper(args: GatherArgs):
+      # Call `precondition_fn`; we assume it may throw an exception.
+      precondition_fn(args)
+      return gather_fn(args)
+
+    return wrapper
+
+  return decorator
+
+
+def _pre_gather_for_scalar_indexing(args: GatherArgs):
+  """Returns True if this call to gather represents scalar indexing into arrays.
+
+  E.g., op[2], op[:, :5, :], jnp.take(op, 0, axis=0).
+  """
+  # TODO(marcvanzee): Add more assumptions here, because this is currently too
+  # permissive.
+  if len(args.start_indices_shape) != 1:
+    raise ValueError("start_indices shape should be 1")
+
+
+@gather_precondition(_pre_gather_for_scalar_indexing)
+def _gather_for_scalar_indexing(args: GatherArgs):
+  """Implements 'scalar indexing into arrays' cases of lax.gather using tf.slice.
+
+  E.g., op[2], op[:, :5, :], jnp.take(op, 0, axis=0).
+  """
+  indices = tf.expand_dims(args.dnums.start_index_map, 1)
+  # lax.gather uses an "index map" which maps `start_indices` to the right axes
+  # in `operand`. Since tf.strided_slice uses a single array for specifying the
+  # start indices, we use a scatter to map the start indices to the right axes.
+  op_shape = jax2tf._eval_shape(args.op_shape)
+  slice_sizes_tf = jax2tf._eval_shape(args.slice_sizes)
+  # TODO(marcvanzee): Consider transposing `operand`, which is probably more
+  # optimization friendly.
+  begin = tf.scatter_nd(indices, args.start_indices, [len(op_shape)])
+  begin = _clip(op_shape, begin, slice_sizes_tf)
+  end = slice_sizes_tf + begin
+
+  # `collapsed_slice_dims` is a tuple of dimensions to collapse, e.g. (0, 2).
+  # `tf.strided_slice` expects a binary mask to specify the shrink axes, i.e.,
+  # if we want to shrink axis 0 and 2, this corresponds to binary mask 101,
+  # which is 5 in decimals. The following line converts the lax representation
+  # to the one used by `tf.strided_slice`.
+  shrink_mask = sum(2**x for x in args.dnums.collapsed_slice_dims)
+  res = tf.strided_slice(args.operand, begin, end, shrink_axis_mask=shrink_mask)
+  # Shape inference doesn't work for tf.strided_slice.
+  res = jax2tf._ensure_tf_shape_if_dynamic(
+      res, jax2tf._aval_to_tf_shape(args.out_aval)
+  )
+  return res
+
+
+def _pre_gather_for_multidim_indexing(args: GatherArgs):
+  """Returns True if this call to gather represents multi-dimensional indexing.
+
+  E.g., jnp.take(op, [[0], [1]], axis=0).
+  Note we currently only support multi-dimensional indexing if the last
+  dimension is 1.
+  """
+  # Handle only the case when tf.gather argument batch_dims=0.
+  # Find axis to match the tf.gather semantics
+  # Let I = len(start_indices_shape)
+  # let O = len(op_shape)
+  # slice_sizes == op_shape[:axis] + (1,) + op_shape[axis+1:]
+  # collapsed_slice_dims == (axis,)
+  # start_index_map == (axis,)
+  # offset_dims == (0, 1, ..., axis - 1, axis + I, ..., O + I - 1)
+  # We added a trailing dimension of size 1
+  op_shape = args.op_shape
+  start_index_map = args.dnums.start_index_map
+  collapsed_slice_dims = args.dnums.collapsed_slice_dims
+  offset_dims = args.dnums.offset_dims
+  if not (len(op_shape) >= 1 and len(start_index_map) == 1 and
+          len(collapsed_slice_dims) == 1 and collapsed_slice_dims[0]
+          == start_index_map[0] and len(offset_dims) == len(op_shape) - 1):
+    raise ValueError("unsupported dimension numbers")
+  # We added a trailing dimension of size 1
+  if not core.definitely_equal(args.start_indices_shape[-1], 1):
+    raise ValueError("start_indices shape[-1] should be 1")
+  # Guess the axis
+  axis = collapsed_slice_dims[0]
+  index_dims = len(args.start_indices_shape) - 1
+  expected_offset_dims = tuple(
+      list(range(axis)) +
+      list(range(axis + index_dims,
+                 len(op_shape) + index_dims - 1)))
+  if offset_dims != expected_offset_dims:
+    raise ValueError("unsupported offset_dims")
+  expected_slice_sizes = op_shape[:axis] + (1,) + op_shape[axis + 1:]  # type: ignore
+  if not core.definitely_equal_shape(args.slice_sizes, expected_slice_sizes):
+    raise ValueError("unsupported slice_sizes")
+
+
+@gather_precondition(_pre_gather_for_multidim_indexing)
+def _gather_for_multidim_indexing(args: GatherArgs):
+  """Implements 'multi-dimensional indexing into arrays' cases of lax.gather using tf.gather.
+
+  E.g., jnp.take(op, [[0], [1]], axis=0).
+  """
+  # Guess the axis.
+  axis = args.dnums.collapsed_slice_dims[0]
+  squeezed_indices = tf.squeeze(args.start_indices, -1)
+  op_shape = jax2tf._eval_shape(args.op_shape)
+  start_indices = _clip((op_shape[axis],), squeezed_indices, (1,))
+  return tf.gather(args.operand, start_indices, axis=axis, batch_dims=0)
+
+
+def _pre_gather_with_batch_dim(args: GatherArgs):
+  """Returns True if this call to gather has non-empty batch dimensions.
+
+  This is for instance triggered when doing jax.vmap(lax.dynamic_slice).
+  """
+  # We assume exactly one batch (and one or more non-batch dimensions).
+  if len(args.batch_dims) != 1:
+    raise ValueError(f"batch_dims is {len(args.batch_dims)} but should be 1")
+
+  # `start_index_map` maps indices in `start_indices` to indices in `operand`.
+  # For simplicity, we currently only consider the case where this mapping is
+  # the identity function, i.e., [2, 3] in `start_indices` maps to
+  # `operand[2, 3]`.
+  if args.dnums.start_index_map != tuple(range(args.start_indices_shape[-1])):
+    raise ValueError("unsupported start_index_map")
+
+  # The batch dims in `start_indices` and `operand` should agree.
+  if not core.definitely_equal(args.op_shape[0], args.start_indices_shape[0]):
+    raise ValueError("Batch dimensions in operand and start_indices don't "
+                     "agree")
+
+
+def _pre_gather_with_batch_dims(args: GatherArgs):
+  """Returns True if this call to gather has non-empty 2D batch dimensions.
+
+  This is for instance triggered when doing
+  jax.vmap(jax.vmap(lax.dynamic_slice)).
+  """
+  if len(args.dnums.collapsed_slice_dims) != 0:
+    # NOTE: this can be relaxed in _gather_with_batch_dims but we might
+    #   also need to re-work the output reshaping
+    raise ValueError("only len(collapsed_slice_dims) == 0 is supported")
+
+  # NOTE: This supports higher dimensions than listed (the highest dimension
+  # in the tests is 3D so it is limited to that, but the implementation is
+  # designed to handle higher dimensions (N-Dimensional)).
+  if len(args.batch_dims) not in [1, 2, 3]:
+    raise ValueError(
+        f"Size of batch_dims is {len(args.batch_dims)} but should be up to 3"
+    )
+
+@gather_precondition(_pre_gather_with_batch_dim)
+def _gather_with_batch_dim(args: GatherArgs):
+  """Implements call to gather with non-empty batch dimensions.
+
+  E.g., when doing `jax.vmap(lax.dynamic_slice).
+  """
+  op_shape = jax2tf._eval_shape(args.op_shape)
+  start_indices = _clip(op_shape, args.start_indices, args.slice_sizes)
+  result = tf.map_fn(
+    lambda idxs: tf.slice(args.operand, begin=idxs, size=args.slice_sizes),
+    start_indices,
+    fn_output_signature=jax2tf._to_tf_dtype(args.operand.dtype)
+  )
+  result = tf.reshape(result, jax2tf._eval_shape(args.out_aval.shape))
+  return result
+
+
+def _gather_generate_indices(shape: tuple[int, ...]):
+  """
+  Returns the indices of the according to `shape`:
+    each element in the output is the index of an element of an array
+    of the provided shape. The result's shape is (math.prod(shape), len(shape))
+
+  For example, given shape (2,2) it returns (0,0),(0,1),(1,0),(1,1)
+  """
+  return tf.reshape(
+      tf.stack(
+          tf.meshgrid(
+              *[tf.range(start=0, limit=x) for x in shape], indexing="ij"
+          ),
+          axis=-1,
+      ),
+      (-1, len(shape)),
+  )
+
+
+@gather_precondition(_pre_gather_with_batch_dims)
+def _gather_with_batch_dims(args: GatherArgs):
+  """Implements call to gather with non-empty 2D batch dimensions."""
+  op_shape = jax2tf._eval_shape(args.op_shape)
+  output_shape = jax2tf._eval_shape(args.out_aval.shape)
+  # Used to map the start_indices w.r.t start_index_map
+  indices = tf.expand_dims(args.dnums.start_index_map, 1)
+
+  # batch_indices is shaped (N,d) where N is the number of slices and d is
+  # the number of batch_dims; batch_indices_size equals to N
+  batch_indices = _gather_generate_indices(
+      tuple(output_shape[i] for i in args.batch_dims)
+  )
+  batch_indices_size = jax2tf._eval_shape(batch_indices.shape)[0]
+  # offset_indices is shaped (K,d) where K is the number of elements in each
+  # slice and d is the number of offset_dims; offset_indices_size equals to K
+  offset_indices = _gather_generate_indices(
+      tuple(output_shape[i] for i in args.dnums.offset_dims)
+  )
+  offset_indices_size = jax2tf._eval_shape(offset_indices.shape)[0]
+
+  # After we compute the result we need to reshape the axes with respect to
+  # the output batch_dims and offset_dims.
+  dim_mask = args.batch_dims + args.dnums.offset_dims
+  mask_output_shape = tuple(output_shape[x] for x in dim_mask)
+
+  def get_scatter_indices(indices, batch_indices_size, size_of_index_map):
+    """Generate the start indices of each slice, which index into the operand."""
+    # Tile indices batch_indices_size times
+    tiled_indices = tf.tile(
+        tf.expand_dims(indices, 0), [batch_indices_size, 1, 1]
+    )
+    # The above tiles need to index the proper element of batch_indices
+    # To do this generate a repeated sequence of numbers
+    temp_batch_indices = tf.repeat(
+        tf.range(start=0, limit=batch_indices_size), size_of_index_map
+    )
+    # Reshape the above sequence so it follows the same shape of tiled_indices
+    temp_batch_indices = tf.reshape(
+        temp_batch_indices, (batch_indices_size, size_of_index_map, 1)
+    )
+    # Now we concatenate to create indices offset by the temp_batch_indices
+    return tf.concat([temp_batch_indices, tiled_indices], axis=-1)
+
+  slice_start_indices = tf.gather_nd(args.start_indices, batch_indices)
+  # TODO: In the case where start_index_map is the identity we can skip this.
+  scatter_indices = get_scatter_indices(
+      indices, batch_indices_size, len(args.dnums.start_index_map)
+  )
+  # We map the scatter_indices w.r.t start_index_map
+  indices_in_operand = tf.scatter_nd(
+        scatter_indices, slice_start_indices, [batch_indices_size, len(op_shape)]
+  )
+
+  # We clip the indices as OOB cases are possible when offsetting past
+  # the operand boundaries
+  clipped_start_indices = _clip(op_shape, indices_in_operand, args.slice_sizes)
+  # Here we need to broadcast clipped_start_indices and add each of the offsets
+  # which will generate a large index tensor of shape (T,d) where T is the
+  # number of slices times the size of each slice (i.e total number of items
+  # across all sices); d is rank(operand)
+  slice_element_indices = tf.add(
+      tf.repeat(clipped_start_indices, offset_indices_size, axis=0),
+      tf.tile(offset_indices, (batch_indices_size, 1)),
+  )
+  results = tf.gather_nd(args.operand, slice_element_indices)
+
+  # Here results comes shaped as (N,1). Because collapsed_slice_dims is 0,
+  # offset_dims is effectviely slice_sizes.
+  # We reshape to mask_output_shape because if we directly reshape to the
+  # output shape and our batch_dims are non-contiguous we will produce the
+  # wrong shape. Reshaping to mask_output_shape gives (...,*slice_sizes),
+  # which we then transpose to permute the axes in the proper way.
+  # Note that if the batch_dims are contiguous this won't change the output.
+  temp = tf.reshape(results, shape=mask_output_shape)
+  return tf.transpose(temp, perm=tf.math.invert_permutation(dim_mask))
+
+def _gather(operand, start_indices, *, dimension_numbers,
+            slice_sizes: core.Shape, indices_are_sorted, unique_indices, mode,
+            fill_value, _in_avals: Sequence[core.ShapedArray],
+            _out_aval: core.ShapedArray):
+  """Tensorflow implementation of gather."""
+  if mode == lax.GatherScatterMode.FILL_OR_DROP:
+    gather_fill_fn = jax2tf._convert_jax_impl(lax_slicing._gather_fill,
+                                              multiple_results=False)
+    return gather_fill_fn(
+        operand, start_indices, dimension_numbers=dimension_numbers,
+        slice_sizes=slice_sizes, unique_indices=unique_indices,
+        indices_are_sorted=indices_are_sorted, fill_value=fill_value,
+        output_shape=_out_aval.shape, _in_avals=_in_avals, _out_aval=_out_aval)
+
+  # TODO(marcvanzee): Check if we need more tests in shape_poly for gather with
+  # enable_xla=False.
+  gather_args = GatherArgs(
+      operand=operand,
+      start_indices=start_indices,
+      dnums=dimension_numbers,
+      slice_sizes=slice_sizes,
+      op_shape=_in_avals[0].shape,
+      start_indices_shape=_in_avals[1].shape,
+      out_aval=_out_aval)
+
+  errors = []
+
+  for gather_fn in [
+      _gather_for_scalar_indexing,
+      _gather_for_multidim_indexing,
+      _gather_with_batch_dim,
+      _gather_with_batch_dims,
+  ]:
+    try:
+      return gather_fn(gather_args)
+    except ValueError as e:
+      errors.append(f"{gather_fn}: {e!r}")
+
+  error_msg = (f"Unsupported arguments for gather: {gather_args}, errors:\n" +
+               "\n".join(errors))
+
+  raise _error("gather", error_msg)
+
+
+tf_impl_no_xla[lax.gather_p] = _gather
+
+
+def _dynamic_slice(operand, *start_indices, slice_sizes: core.Shape,
+                   _in_avals: Sequence[core.ShapedArray],
+                   _out_aval: core.ShapedArray):
+  start_indices = tf.stack(start_indices)
+  slice_sizes_tf = jax2tf._eval_shape(slice_sizes)
+
+  operand_shape = jax2tf._eval_shape(_in_avals[0].shape)
+  start_indices = _clip(operand_shape, start_indices, slice_sizes_tf)
+  return tf.slice(operand, start_indices, size=slice_sizes_tf)
+
+
+tf_impl_no_xla[lax.dynamic_slice_p] = _dynamic_slice
+
+
+def _dynamic_update_slice(operand, update, *start_indices,
+                          _in_avals: Sequence[core.ShapedArray],
+                          _out_aval: core.ShapedArray):
+  start_indices = tf.stack(start_indices)
+
+  op_shape = jax2tf._eval_shape(_in_avals[0].shape)
+  op_size = tf.size(operand)
+  update_shape_tf = jax2tf._eval_shape(_in_avals[1].shape)
+
+  start_indices = _clip(op_shape, start_indices, update_shape_tf)
+  end_indices = tf.add(start_indices, update_shape_tf)
+
+  # Get the cells to update in `operand` as an array of ids.
+  id_tensor = tf.reshape(tf.range(op_size), op_shape)
+  scattered_indices = tf.strided_slice(id_tensor, start_indices, end_indices)
+
+  # Create an array containing updates at scattered_indices and zeros otherwise.
+  flat_indices = tf.expand_dims(tf.nest.flatten(scattered_indices), -1)
+  flat_update = tf.nest.flatten(update)
+  update = tf.scatter_nd(flat_indices, flat_update, (op_size,))
+  update = tf.reshape(update, op_shape)
+
+  # Create a bool mask that is True only where `operand` should be updated.
+  update_mask = tf.ones_like(flat_update, dtype=tf.bool)
+  update_mask = tf.scatter_nd(flat_indices, update_mask, (op_size,))
+  update_mask = tf.reshape(update_mask, op_shape)
+
+  # Use the mask to only update `operand` with `update`.
+  return tf.where(update_mask, update, operand)
+
+
+tf_impl_no_xla[lax.dynamic_update_slice_p] = _dynamic_update_slice
+
+
+def shift_axes_forward(operand,
+                       axes: tuple[int, ...],
+                       inverse: bool = False,
+                       forward: bool = True):
+  """Shifts the tuple of axes to the front of an array"""
+  other_axes = tuple(i for i in range(len(operand.shape)) if i not in axes)
+  fwd_order = axes + other_axes if forward else other_axes + axes
+  order = fwd_order if not inverse else _invert_permutation(fwd_order)
+  return tf.transpose(operand, order)
+
+def convert_scatter_jax_to_tf(update_op, unsorted_segment_op=None):
+
+  def _sparse_scatter(operand, scatter_indices, updates, unique_indices, mode,
+                      _in_avals: Sequence[core.ShapedArray],
+                      _out_aval: core.ShapedArray):
+    """Implementation of scatter specialised to indexing from the front axes.
+
+    This covers unique indices and non-unique indices of single depth.
+    Note on unique indices: `tf.tensor_scatter_nd_update` interprets indices
+    thusly: every axis except the final one encodes a batch dimension, the final
+    axis encoding the actual indices to scatter in to. It enforces, at least
+    one, batch dimension so we add an empty dimension to indices and updates if
+    lacking.
+
+    Note on non-unique indices: There is no tf op for non-single depth indexing,
+    but if indexing is single depth, this can be viewed as a segment op.
+    """
+    # Infer unique indices from lack of batch dimension
+    unique_indices = unique_indices or (len(scatter_indices.shape) == 1)
+    if unique_indices:
+      suboperand = tf.gather_nd(operand, scatter_indices)
+      updated_suboperand = update_op(suboperand, updates)
+      # add a batch dim if none exist
+      if len(scatter_indices.shape) == 1:
+        scatter_indices = scatter_indices[None]
+        updated_suboperand = updated_suboperand[None]
+      y = tf.tensor_scatter_nd_update(operand, scatter_indices, updated_suboperand)
+    else:
+      if (scatter_indices.shape[-1] == 1) and unsorted_segment_op:
+        # If only indexing into the first dimension, it's a segment op
+        operand_update = unsorted_segment_op(updates,
+                                             tf.squeeze(scatter_indices, -1),
+                                             operand.shape[0])
+        y = update_op(operand, operand_update)
+      else:
+        raise _scatter_error(
+            "Scatter only supports non-unique "
+            "indices with indexing into only one dimension for (add, mul, min, "
+            "max)")
+    return y
+
+  def sparse_scatter(operand, scatter_indices, updates, update_jaxpr,
+                     update_consts, dimension_numbers, indices_are_sorted: bool,
+                     unique_indices: bool, mode,
+                     _in_avals: Sequence[core.ShapedArray],
+                     _out_aval: core.ShapedArray):
+    """
+    Wrapper around the scatter function.
+    The underlying tf ops `tf.tensor_scatter_nd_update` and
+    `tf.math.unsorted_segment_*` index from the front dimensions.
+    `tf.math.unsorted_segment_*` indexes to a depth 1 from the front.
+    `tf.tensor_scatter_nd_update` indexes from the front dimensions onwards,
+    with no ability to skip a dimension. This function shifts the axes to be
+    indexed to the front then calls a front-specific implementation, then
+    inverse-shifts the output.
+
+    scatter_dims_to_operand_dims: dimensions which the scatter indexes in to.
+      We shift these to the front to match tf syntax. All other dims are batch
+    update_window_dims: dimensions which are not batch dimensions. We shift
+      these to the back as the remaining dimensions are batch dimensions.
+    """
+    del update_jaxpr, update_consts, indices_are_sorted  # Unused arguments
+
+    update_window_dims = dimension_numbers.update_window_dims
+    inserted_window_dims = dimension_numbers.inserted_window_dims
+    scatter_to_operand_dims = dimension_numbers.scatter_dims_to_operand_dims
+
+    dtype = operand.dtype  # assume updates has same dtype as operand
+    if dtype in [tf.bool, tf.complex64]:
+      raise _scatter_error(f"Scatter does not support operands of type {dtype}")
+
+    if inserted_window_dims != scatter_to_operand_dims:
+      raise _scatter_error("Complex scatters are not supported")
+
+    if (mode != lax.GatherScatterMode.FILL_OR_DROP and
+        mode != lax.GatherScatterMode.PROMISE_IN_BOUNDS):
+      # The OOB behavior for tf.scatter is as follows:
+      # - When running in eager or graph mode, it throws an error.
+      #   TODO(marcvanzee): Fix this case by removing the OOB indices.
+      # - When running in compile mode, the OOB indices are dropped, which is
+      #   the same behavior as FILL_OR_DROP and PROMISE_IN_BOUNDS.
+      # To ensure correctness, we disallow CLIP mode for now.
+      raise _scatter_error("Only scatter modes `FILL_OR_DROP` and "
+                           "`PROMISE_IN_BOUNDS` are supported.")
+
+    # Shift axes to the front to match tf syntax, inverse afterwards
+    fwd = partial(shift_axes_forward, axes=scatter_to_operand_dims)
+    inv = partial(fwd, inverse=True)
+
+    # Shift update value axes to the back, so batch are at the front
+    updates_shifted = shift_axes_forward(
+        updates, axes=update_window_dims, forward=False)
+
+    return inv(
+        _sparse_scatter(
+            fwd(operand), scatter_indices, updates_shifted, unique_indices,
+            mode, _in_avals, _out_aval))
+  return sparse_scatter
+
+
+tf_impl_no_xla[lax.scatter_p] = convert_scatter_jax_to_tf(
+    lambda x, y: y)  # just replace with the update
+tf_impl_no_xla[lax.scatter_add_p] = convert_scatter_jax_to_tf(tf.add,      tf.math.unsorted_segment_sum)
+tf_impl_no_xla[lax.scatter_mul_p] = convert_scatter_jax_to_tf(tf.multiply, tf.math.unsorted_segment_prod)
+tf_impl_no_xla[lax.scatter_min_p] = convert_scatter_jax_to_tf(tf.minimum,  tf.math.unsorted_segment_min)
+tf_impl_no_xla[lax.scatter_max_p] = convert_scatter_jax_to_tf(tf.maximum,  tf.math.unsorted_segment_max)
+
+tf_impl_no_xla[lax.sort_p] = _unimplemented("sort")
+
+tf_impl_no_xla[lax.reduce_precision_p] = _unimplemented("reduce_precision")

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/tests/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2020 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/tests/back_compat_tf_test.py

@@ -0,0 +1,154 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for backwards compatibility of custom calls involving TensorFlow.
+
+See the back_compat_test_util module docstring for how to setup and update
+these tests.
+"""
+
+from __future__ import annotations
+
+import base64
+from collections.abc import Sequence
+import io
+import os
+import tarfile
+from typing import Callable, Optional
+
+from absl.testing import absltest
+import jax
+from jax._src import test_util as jtu
+from jax._src.internal_test_util import export_back_compat_test_util as bctu
+from jax._src.lib import xla_extension
+from jax.experimental import jax2tf
+from jax.experimental.jax2tf.tests.back_compat_testdata import tf_call_tf_function
+import jax.numpy as jnp
+import tensorflow as tf
+
+
+jax.config.parse_flags_with_absl()
+
+
+def serialize_directory(directory_path):
+  """Seriliaze the directory as a string."""
+  tar_buffer = io.BytesIO()
+  with tarfile.open(fileobj=tar_buffer, mode="w") as tar:
+    tar.add(directory_path, arcname=os.path.basename(directory_path))
+
+  # Convert the binary data to a base64-encoded string
+  serialized_string = base64.b64encode(tar_buffer.getvalue())
+  return serialized_string
+
+
+def deserialize_directory(serialized_string, output_directory):
+  """Deserialize the string to the directory."""
+  # Convert the base64-encoded string back to binary data
+  tar_data = base64.b64decode(serialized_string)
+
+  # Extract the tar archive to the output directory
+  with tarfile.open(fileobj=io.BytesIO(tar_data), mode="r") as tar:
+    tar.extractall(output_directory)
+
+
+class CompatTensoflowTest(bctu.CompatTestBase):
+  """Compatibility tests that use TF.
+
+  Uses tf.Graph to serialize and run the functions; expects that `func`
+  contains a `jax2tf.call_tf` and uses `jax2tf.convert` to generate a
+  `tf.Graph` containing a XlaCallModule with the actual MLIR module.
+  """
+
+  def run_current(self, func: Callable, data: bctu.CompatTestData):
+    # Here we use tf.saved_model and provide  string serialize/deserialize methods
+    # for the whole directory.
+    @tf.function(autograph=False, jit_compile=True)
+    def tf_func(the_input):  # Use recognizable names for input and result
+      res = jax2tf.convert(func, native_serialization=True)(the_input)
+      return tf.identity(res, name="the_result")
+
+    self.tf_func = tf_func
+    return tf_func(*data.inputs)
+
+  def serialize(
+      self,
+      func: Callable,
+      data: bctu.CompatTestData,
+      polymorphic_shapes: Sequence[str] | None = None,
+      allow_unstable_custom_call_targets: Sequence[str] = (),
+  ):
+    # We serialize as a tf.Graph
+    assert len(data.inputs) == 1  # We only support a single input now
+    tf_graph = self.tf_func.get_concrete_function(*data.inputs).graph
+    for op in tf_graph.get_operations():
+      if op.type == "XlaCallModule":
+        serialized_module = op.get_attr("module")
+        module_str = xla_extension.mlir.deserialize_portable_artifact(
+            serialized_module
+        )
+        module_version = op.get_attr("version")
+        break
+    else:
+      raise ValueError("Cannot find an XlaCallModule")
+    tf_graph_def = tf_graph.as_graph_def()
+    # module_str is just for human readability, add both the MLIR module
+    # and the tf.Graph
+    module_str = (
+        "# First the MLIR module:\n"
+        + module_str
+        + "\n# Then the tf.Graph:\n"
+        + str(tf_graph_def)
+    )
+    # serialized = tf_graph_def.SerializeToString()
+    module = tf.Module()
+    module.call = self.tf_func.get_concrete_function(*data.inputs)
+    root_dir = self.create_tempdir()
+    saved_model_dir = os.path.join(root_dir, "saved_model")
+    os.mkdir(saved_model_dir)
+    tf.saved_model.save(
+        module,
+        saved_model_dir,
+        options=tf.saved_model.SaveOptions(experimental_custom_gradients=False),
+    )
+    serialized =  serialize_directory(saved_model_dir)
+    nr_devices = 1
+    return serialized, module_str, module_version, nr_devices
+
+  def run_serialized(
+      self,
+      data: bctu.CompatTestData,
+      polymorphic_shapes: Sequence[str] | None = None,
+  ):
+    root_dir = self.create_tempdir()
+    deserialize_directory(data.mlir_module_serialized, root_dir)
+    saved_model_dir = os.path.join(root_dir, "saved_model")
+    loaded_model = tf.saved_model.load(saved_model_dir)
+    return (loaded_model.call(*data.inputs).numpy(),)
+
+  def test_tf_call_tf_function(self):
+    # A custom call tf.call_tf_function is generated when we lower call_tf
+    # with the call_tf_graph=True option.
+    def func(x):
+      def func_tf(x):
+        return tf.math.sin(x)
+
+      return jnp.cos(
+          jax2tf.call_tf(func_tf, output_shape_dtype=x, call_tf_graph=True)(x)
+      )
+
+    data = self.load_testdata(tf_call_tf_function.data_2023_07_29)
+    self.run_one_test(func, data)
+
+
+if __name__ == "__main__":
+  absltest.main(testLoader=jtu.JaxTestLoader())

external/alphageometry/.venv-ag/Lib/site-packages/jax/experimental/jax2tf/tests/call_tf_test.py

@@ -0,0 +1,1821 @@
+# Copyright 2020 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for call_tf."""
+
+import contextlib
+from functools import partial
+import os
+from typing import Callable
+import unittest
+
+from absl import logging
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+from jax import dlpack
+from jax import dtypes
+from jax import lax
+from jax import numpy as jnp
+from jax._src import config
+from jax._src import test_util as jtu
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo
+from jax.experimental import export
+from jax.experimental import jax2tf
+from jax.experimental.jax2tf.tests import tf_test_util
+import numpy as np
+
+try:
+  import tensorflow as tf
+except ImportError:
+  tf = None
+
+jax.config.parse_flags_with_absl()
+
+
+def _maybe_jit(with_jit: bool, func: Callable) -> Callable:
+  if with_jit:
+    return jax.jit(func)
+  else:
+    return func
+
+def _maybe_tf_jit(with_jit: bool, func: Callable) -> Callable:
+  if with_jit:
+    return tf.function(func, autograph=False, jit_compile=True)
+  else:
+    return func
+
+def _named_test(**kwargs):
+  return dict(kwargs,
+              testcase_name = "_".join([f"{k}={kwargs[k]}" for k in sorted(kwargs.keys())]))
+
+_parameterized_jit = parameterized.named_parameters(
+    _named_test(with_jit=with_jit)
+    for with_jit in [True, False])
+
+_call_tf_non_compilable_error = "Error compiling TensorFlow function"
+_call_tf_dynamic_shape_error = "call_tf cannot call functions whose output has dynamic shape"
+
+class CallTfTest(tf_test_util.JaxToTfTestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    # One TF device of each device_type
+    cls.tf_devices = []
+    for tf_device in tf.config.list_logical_devices():
+      if tf_device.device_type == "TPU_SYSTEM":
+        continue  # A virtual device
+      if all(tf_device.device_type != d.device_type for d in cls.tf_devices):
+        cls.tf_devices.append(tf_device)
+
+    super().setUpClass()
+
+  def setUp(self):
+    if tf is None:
+      raise unittest.SkipTest("Test requires tensorflow")
+    # TODO(b/171320191): this line works around a missing context initialization
+    # bug in TensorFlow.
+    _ = tf.add(1, 1)
+    super().setUp()
+
+  @_parameterized_jit
+  def test_eval_scalar_arg(self, with_jit=True):
+    def f_tf(x):
+      return tf.math.sin(x)
+    x = 3.
+    res = _maybe_jit(with_jit, jax2tf.call_tf(f_tf))(x)
+    self.assertAllClose(jnp.sin(x), res)
+
+  @_parameterized_jit
+  def test_eval_scalar_res(self, with_jit=True):
+    x = 3.
+    res = _maybe_jit(with_jit, jax2tf.call_tf(lambda x: 4.))(x)
+    self.assertAllClose(4., res, check_dtypes=False)
+
+  @_parameterized_jit
+  def test_eval_numpy_arg(self, with_jit=True):
+    x = np.ones((2, 3), dtype=np.float32)
+    res = _maybe_jit(with_jit, jax2tf.call_tf(tf.math.sin))(x)
+    self.assertAllClose(jnp.sin(x), res)
+
+  @_parameterized_jit
+  def test_eval_numpy_res(self, with_jit=False):
+    x = np.ones((2, 3))
+    res = _maybe_jit(with_jit, jax2tf.call_tf(lambda _: x))(x)
+    self.assertAllClose(x, res)
+
+  @_parameterized_jit
+  def test_eval_devicearray_arg(self, with_jit=False):
+    x = jnp.ones((2, 3), dtype=np.float32)
+    res = _maybe_jit(with_jit, jax2tf.call_tf(tf.math.sin))(x)
+    self.assertAllClose(jnp.sin(x), res)
+
+    x = jnp.array(3.0, dtype=jnp.bfloat16)
+    res = jax2tf.call_tf(lambda x: x)(x)
+    self.assertAllClose(x, res)
+    # bfloat16 scalar will create a copy.
+    with self.assertRaises(AssertionError):
+      self.assertTrue(np.shares_memory(x, res))
+
+  @_parameterized_jit
+  def test_eval_pytree(self, with_jit=True):
+
+    def fun_tf(x: dict, y: tuple) -> tuple:
+      return (x["first"] * x["second"], y[0] + y[1])
+
+    x = dict(first=np.float32(3.), second=np.float32(4.))
+    y = (np.float64(5.), np.float64(6.))
+    fun_jax = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))
+    res = fun_jax(x, y)
+    self.assertAllClose((np.float32(12.), np.float64(11.)), res)
+
+  def test_result_tuple(self):
+    x1 = np.ones(3, dtype=np.int32)
+    x2 = np.ones(5, dtype=np.float32)
+    def fun_tf():
+      return tf.tuple([x1, x2])
+
+    fun_jax = jax.jit(jax2tf.call_tf(fun_tf))
+    res = fun_jax()
+    self.assertAllClose(res, (x1, x2))
+
+  def test_error_non_compilable_strings(self):
+    # Check that in op-by-op we call a function in eager mode.
+    def f_tf_non_compilable(x):
+      return tf.strings.length(tf.strings.format("Hello {}!", [x]))
+
+    f_jax = jax2tf.call_tf(f_tf_non_compilable)
+    x = np.float32(0.7)
+    self.assertAllClose(f_tf_non_compilable(x).numpy(), f_jax(x))
+    with self.assertRaisesRegex(ValueError,
+                                _call_tf_non_compilable_error):
+      jax.jit(f_jax)(x)
+
+    with self.assertRaisesRegex(ValueError,
+                                _call_tf_non_compilable_error):
+      lax.cond(True, lambda x: f_jax(x), lambda x: f_jax(x), x)
+
+  def test_error_non_compilable_dynamic_shape(self):
+    # Check that in op-by-op we call a function in eager mode.
+    def f_tf_non_compilable(x):
+      return tf.cond(x[0], lambda: x[1:], lambda: x)
+
+    f_jax = jax2tf.call_tf(f_tf_non_compilable)
+    x = np.array([True, False], dtype=np.bool_)
+    self.assertAllClose(f_tf_non_compilable(x), f_jax(x))  # Works in eager mode
+
+    with self.assertRaisesRegex(ValueError, _call_tf_dynamic_shape_error):
+      jax.jit(f_jax)(x)
+
+  def test_error_bad_result_tensorarray(self):
+    # Call a function that returns a tf.TensorArray. This should be detected
+    # early on. If we don't the function is actually compilable but returns
+    # a tuple instead of a single result.
+    def fun_tf():
+      ta = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
+      ta = ta.unstack([0, 1, 2, 3, 4])
+      return ta
+
+    with self.assertRaisesRegex(ValueError,
+                                "The called TF function returns a result that is not convertible to JAX"):
+      fun_jax = jax.jit(jax2tf.call_tf(fun_tf))
+      fun_jax()
+
+  def test_error_bad_result_string(self):
+    def fun_tf():
+      return tf.constant("foo")
+
+    # Now under jit, should fail because the function is not compilable
+    with self.assertRaisesRegex(ValueError,
+                                "The called TF function returns a result that is not convertible to JAX"):
+      fun_jax = jax.jit(jax2tf.call_tf(fun_tf))
+      fun_jax()
+
+  @_parameterized_jit
+  def test_control_flow(self, with_jit=True):
+
+    def times_5_tf(x):
+      # Multiply x * 5 using a loop
+      c = lambda i, acc: tf.less(i, 5)
+      b = lambda i, acc: (tf.add(i, 1), tf.add(acc, x))
+      _, acc = tf.while_loop(c, b, [tf.constant(0), tf.constant(0.)])
+      return acc
+
+    def fun_jax(x):
+      # Calls times_5_tf 3 times in a loop
+      def body(_, acc):
+        return jax2tf.call_tf(times_5_tf)(acc)
+
+      return lax.fori_loop(0, 3, body, x)
+
+    x = np.float32(3.)
+    res = _maybe_jit(with_jit, fun_jax)(x)
+    self.assertAllClose(np.float32(x * 5 * 5 * 5), res)
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name=f"_{dtype.__name__}{'_jit' if with_jit else ''}",
+          dtype=dtype,
+          with_jit=with_jit)
+      for dtype in set(jtu.dtypes.all) - {np.bool_}
+      for with_jit in [True, False])
+  def test_dtypes(self, dtype=np.int32, with_jit=True):
+
+    def fun_tf(x):
+      # AddV2 supports more types
+      return tf.raw_ops.AddV2(x=x, y=tf.constant(3, dtype=dtype))
+
+    def fun_jax(x):
+      return jax2tf.call_tf(fun_tf)(x) + x
+
+    x = np.ones((3,), dtype=dtype)
+    res = _maybe_jit(with_jit, fun_jax)(x)
+    self.assertAllClose(dtype(2 * x + 3), res)
+
+  @_parameterized_jit
+  def test_bool(self, with_jit=False):
+
+    def fun_tf(x, y):
+      return tf.math.logical_and(x, y)
+
+    x = np.array([True, False, True, False], dtype=np.bool_)
+    y = np.array([True, True, False, False], dtype=np.bool_)
+    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x, y)
+    self.assertAllClose(
+        np.array([True, False, False, False], dtype=np.bool_), res)
+
+  @_parameterized_jit
+  def test_x64_input(self, with_jit=True):
+    def f_tf(x):
+      return tf.math.sin(x)
+
+    x = 5.  # TF interprets this as f64
+    res_call_tf = _maybe_jit(with_jit, jax2tf.call_tf(f_tf))(x)
+    res_jax = jnp.sin(x)
+    self.assertAllClose(res_call_tf, res_jax)
+
+  @_parameterized_jit
+  def test_x64_output(self, with_jit=True):
+    def f_tf(x):
+      return (tf.constant(3., tf.float64), x)
+
+    x = np.float32(5.)
+    res_call_tf = _maybe_jit(with_jit, jax2tf.call_tf(f_tf))(x)
+    res_jax = (3., x)
+    self.assertAllClose(res_call_tf, res_jax)
+
+    res_call_tf_jit = jax.jit(jax2tf.call_tf(f_tf))(x)
+    self.assertAllClose(res_call_tf_jit, res_jax)
+
+  @_parameterized_jit
+  def test_with_var_read(self, with_jit=True):
+    # The variable is placed on the default TF device.
+    outer_var_array = np.array([3., 4.], dtype=np.float32)
+    outer_var = tf.Variable(outer_var_array)
+
+    def fun_tf(x):
+      return x * outer_var + 1.
+
+    x = np.array([2., 5.,], dtype=np.float32)
+    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
+    self.assertAllClose(x * outer_var_array + 1., res, check_dtypes=False)
+
+  @_parameterized_jit
+  def test_with_var_read_x64(self, with_jit=True):
+    outer_var_array = np.array([3., 4.], dtype=np.float64)
+    outer_var = tf.Variable(outer_var_array)
+
+    def fun_tf(x):
+      return x * tf.cast(outer_var, x.dtype) + 1.
+
+    x = np.array([2., 5.,], dtype=np.float32)
+    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
+    self.assertAllClose(x * outer_var_array + 1., res, check_dtypes=False)
+
+  def test_with_var_different_shape(self):
+    # See https://github.com/google/jax/issues/6050
+    v = tf.Variable((4., 2.), dtype=tf.float32)
+
+    def tf_func(x):
+      return v + x
+    x = np.float32(123.)
+    tf_out = tf_func(x)
+
+    jax_func = jax.jit(jax2tf.call_tf(tf_func))
+    jax_out = jax_func(x)
+
+    self.assertAllClose(tf_out, jax_out, check_dtypes=False)
+
+  @_parameterized_jit
+  def test_with_var_write_error(self, with_jit=True):
+    if with_jit:
+      raise unittest.SkipTest("variable writes not yet working")
+    outer_var = tf.Variable(3., dtype=np.float32)
+
+    def fun_tf(x):
+      outer_var.assign(tf.constant(4.))
+      return x * outer_var + 1.
+
+    x = np.float32(2.)
+    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
+    self.assertAllClose(x * 4. + 1, res, check_dtypes=False)
+
+  @_parameterized_jit
+  def test_with_tensor_capture(self, with_jit=True):
+    outer_tensor = tf.constant(3., dtype=np.float32)
+
+    def fun_tf(x):
+      return x * outer_tensor + 1.
+
+    x = np.float32(2.)
+    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
+    self.assertAllClose(x * 3. + 1., res, check_dtypes=False)
+
+  @_parameterized_jit
+  def test_with_tensor_capture_x64(self, with_jit=True):
+    outer_tensor = tf.constant(3., dtype=np.float64)
+
+    def fun_tf(x):
+      return x * tf.cast(outer_tensor * 3.14, tf.float32) + 1.
+
+    x = np.float32(2.)
+    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
+    self.assertAllClose(x * 3. * 3.14 + 1., res, check_dtypes=False)
+
+  @_parameterized_jit
+  def test_with_value_capture(self, with_jit=True):
+    outer_val = np.array(3., dtype=np.float32)
+
+    def fun_tf(x):
+      return x * outer_val + 1.
+
+    x = np.float32(2.)
+    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
+    self.assertAllClose(x * 3. + 1., res, check_dtypes=False)
+
+  @_parameterized_jit
+  def test_with_multiple_capture(self, with_jit=True):
+    if jtu.test_device_matches(["gpu"]):
+      raise unittest.SkipTest("Test fails on GPU")
+    v2 = tf.Variable(2., dtype=np.float32)
+    v3 = tf.Variable(3., dtype=np.float32)
+    t4 = tf.constant(4., dtype=np.float32)
+    t5 = tf.constant(5., dtype=np.float32)
+
+    def fun_tf(x):
+      return (x * v3 + t4 + v2) * v3 + t5
+
+    x = np.float32(2.)
+    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
+    self.assertAllClose((x * 3. + 4. + 2.) * 3. + 5., res, check_dtypes=False)
+
+  @_parameterized_jit
+  def test_grad(self, with_jit=False):
+    x = np.float32(3.)
+    res = _maybe_jit(with_jit, jax.grad(jax2tf.call_tf(tf.math.sin)))(x)
+    self.assertAllClose(np.cos(x), res)
+
+  @_parameterized_jit
+  def test_grad_pytree(self, with_jit=False):
+
+    def fun_tf(x: dict, y: tuple) -> tuple:
+      return x["first"] * x["second"] + 3. * y[0] + 4. * y[1]
+
+    x = dict(first=np.float32(3.), second=np.float32(4.))
+    y = (np.float32(5.), np.float32(6.))
+    grad_x = _maybe_jit(with_jit, jax.grad(jax2tf.call_tf(fun_tf)))(x, y)
+    self.assertAllClose(
+        dict(first=np.float32(4.), second=np.float32(3.)), grad_x)
+
+  def test_grad_nested(self):
+    # We embed the call_tf function in a larger function whose gradient we take
+    # It is relevant here that the cotangents flowing through the call_tf
+    # function are not scalars.
+
+    b = np.array([[11., 12., 13.], [21., 22., 23.]], dtype=np.float32)  # [2, 3]
+    c = np.array([[31., 32.], [41., 42.], [51., 52.], [61., 62.]], dtype=np.float32)  # [4, 2]
+    x_dict = dict(b=b, c=c)  # b:[2, 3], c=[4, 2]
+    # res: dict(r:[4, 3], s:[4, 2])
+    def f_tf(x_dict):
+      return dict(r=tf.matmul(x_dict["c"], x_dict["b"]), s=7. * x_dict["c"])
+
+    @jax.jit  # To recognize it in jaxpr
+    def f_jax(x_dict):
+      return dict(r=jnp.matmul(x_dict["c"], x_dict["b"]), s=7. * x_dict["c"])
+
+    def loss(functional, x_dict):
+      prediction = functional(x_dict)  # r:[4, 3], s:[4, 2]
+      weights = np.array([1., 2., 3., 4.], dtype=np.float32)  # [4]
+      weighted_pred = jnp.matmul(weights, prediction["r"])  # [3]
+      return jnp.sum(weighted_pred) + 4. * jnp.sum(prediction["s"])
+
+    g_fun_with_tf = jax.grad(partial(loss, jax2tf.call_tf(f_tf)))
+    g_fun_with_jax = jax.grad(partial(loss, f_jax))
+
+    g_tf = g_fun_with_tf(x_dict)
+    g_jax = g_fun_with_jax(x_dict)
+    self.assertAllClose(g_jax, g_tf)
+
+  def test_grad_int_argument(self):
+    # Similar to https://github.com/google/jax/issues/6975
+    # state is a pytree that contains an integer and a boolean.
+    # The function returns an integer and a boolean.
+    def f(param, state, x):
+      return param * x, state
+
+    param = np.array([0.7, 0.9], dtype=np.float32)
+    state = dict(array=np.float32(1.), counter=7, truth=True)
+    x = np.float32(3.)
+
+    # tf.function is important, without it the bug does not appear
+    f_call_tf = jax2tf.call_tf(f)
+    g_call_tf = jax.grad(lambda *args: jnp.sum(f_call_tf(*args)[0]))(param, state, x)
+    g = jax.grad(lambda *args: jnp.sum(f(*args)[0]))(param, state, x)
+    self.assertAllClose(g_call_tf, g)
+
+  def test_grad_int_argument_unused(self):
+    batch_size = 5
+    inputs = np.ones((batch_size, 3), dtype=np.float32)
+    rng = np.array([1, 2], dtype=np.uint32)
+    params = np.float32(.5)
+
+    # rng is integer, unused
+    def jax_model(params, rng, inputs):
+      return jnp.ones([batch_size, 2], dtype=jnp.float32)
+
+    tf_model = jax2tf.convert(jax_model, with_gradient=True)
+
+    def _loss_fn(inference_fn, params, rng, inputs):
+      prediction = inference_fn(params, rng, inputs)
+      return jnp.mean(prediction)
+
+    jax_loss_fn = partial(_loss_fn, jax_model)
+    jax_grad = jax.grad(jax_loss_fn)(params, rng, inputs)
+
+    paramsv = tf.Variable(params)
+    with tf.GradientTape() as tape:
+      tf_prediction = tf_model(paramsv, rng, inputs)
+      tf_loss = tf.reduce_mean(tf_prediction)
+
+      tf_grad = tape.gradient(tf_loss, paramsv)
+    self.assertAllClose(jax_grad, tf_grad.numpy())
+
+    call_tf_loss_fn = partial(_loss_fn, jax2tf.call_tf(tf_model))
+    call_tf_grad = jax.grad(call_tf_loss_fn)(params, rng, inputs)
+    self.assertAllClose(jax_grad, call_tf_grad)
+
+  def test_grad_with_float0_result(self):
+    # Gradient over integer-argument functions, with float0 result
+    def f_jax(x, y):  # x is an int, y is a float; res is a (int, float)
+      return (2 * x, 2 * x + y * y)
+    def f_tf(x, y):
+      # TF needs explicit casts
+      return (2 * x, tf.cast(2 * x, dtype=y.dtype) + y * y)
+
+    def wrapper(functional, x, y):  # x: i32
+      return jnp.sum(2. * functional(3 * x, 4. * y)[1])
+
+    grad_g = jax.grad(partial(wrapper, f_jax),
+                      allow_int=True, argnums=(0, 1))
+    grad_g_call_tf = jax.grad(partial(wrapper, jax2tf.call_tf(f_tf)),
+                              allow_int=True, argnums=(0, 1))
+
+    x = np.int32(2)
+    y = np.float32(3.)
+    g_jax = grad_g(x, y)
+    g_call_tf = grad_g_call_tf(x, y)
+    self.assertEqual(g_jax[0].dtype, dtypes.float0)
+    self.assertEqual(g_call_tf[0].dtype, dtypes.float0)
+    self.assertAllClose(g_jax[1], g_call_tf[1])
+
+  @_parameterized_jit
+  def test_grad_custom(self, with_jit=False):
+
+    @tf.custom_gradient
+    def func_square_tf(x):
+      # Like x ** 2, but with custom grad 3. * x
+      def grad(dy, variables=None):
+        # dy, = dys
+        return 3. * x * dy,
+
+      return x * x, grad
+
+    x = np.float32(4.)
+    grad_x = _maybe_jit(with_jit, jax.grad(jax2tf.call_tf(func_square_tf)))(x)
+    self.assertAllClose(np.float32(3.) * x, grad_x)
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name=f"_{degree=}{'_jit' if with_jit else ''}",
+          degree=degree,
+          with_jit=with_jit)
+      for degree in [1, 2, 3, 4]
+      for with_jit in [True, False])
+  def test_higher_order_grad(self, degree=2, with_jit=False):
+
+    def fun_tf(x):
+      return 2. * x * x * x
+
+    def fun_jax(x):
+      return 3. * _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
+
+    def fun_jax_pure(x):
+      return 3. * fun_tf(x)
+
+    grad_jax = fun_jax
+    grad_jax_pure = fun_jax_pure
+    for _ in range(degree):
+      grad_jax = jax.grad(grad_jax)
+      grad_jax_pure = jax.grad(grad_jax_pure)
+
+    res_jax = grad_jax(np.float32(5.))
+    logging.info("Grad of %s degree is %s", degree, res_jax)
+    self.assertAllClose(res_jax, grad_jax_pure(np.float32(5.)))
+
+  def test_pmap(self):
+    logging.info("Running test_pmap on %s devices", jax.local_device_count())
+
+    def plus_2_tf(x):
+      return tf.math.add(2., x)
+
+    def fun_jax(x):
+      return np.float32(3.) * jax2tf.call_tf(plus_2_tf)(x)
+
+    x = np.arange(jax.local_device_count(), dtype=np.float32)
+    res = jax.pmap(fun_jax)(x)
+    self.assertAllClose(np.float32(3. * (x + 2)), res)
+
+  def test_function_compile_time_constant_inputs(self):
+    # Call a function for which shape inference does not give an output
+    # shape.
+    x = np.array([1, 2, 3], dtype=np.int32)
+    def fun_tf(x):  # x:i32[3]
+      # Indexing with a dynamic slice makes the TF shape inference return
+      # a partially known shape.
+      end_idx = x[1]
+      res = x[0:end_idx]
+      return res
+
+    # Call in eager mode. Should work!
+    res1 = jax2tf.call_tf(fun_tf)(x)
+    self.assertAllClose(x[0:x[1]], res1)
+
+    # Now under jit, should fail because the function is not compilable
+    with self.assertRaisesRegex(ValueError, _call_tf_dynamic_shape_error):
+      fun_jax = jax.jit(jax2tf.call_tf(fun_tf))
+      fun_jax(x)
+
+  def test_experimental_get_compiler_ir_design_doc(self):
+    # Not a test of call_tf, but more of how experimental_get_compiler_ir works.
+    # Examples are from the design doc.
+
+    # Constant slice. This is the common case.
+    x = np.zeros((10,), dtype=np.int32)
+
+    def fun_tf(x):
+      begin = 0
+      return x[begin:5]
+
+    hlo = tf.function(fun_tf, jit_compile=True, autograph=False).experimental_get_compiler_ir(x)()
+    self.assertIn("(arg0.1: s32[10]) -> s32[5]", hlo)
+
+    # Non-constant slice, but compile-time constant depending only on values.
+    x = np.zeros((10,), dtype=np.int32)
+
+    # Non-constant slice, but compile-time constant depending only on shapes.
+    x = np.zeros((10,), dtype=np.int32)
+
+    def fun_tf(x):
+      begin = tf.shape(x)[0] - 2  # begin is a compile-time constant, even if x is not
+      return x[begin:]
+
+    hlo = tf.function(fun_tf, jit_compile=True, autograph=False).experimental_get_compiler_ir(x)()
+    self.assertIn("(arg0.1: s32[10]) -> s32[2]", hlo)
+
+    # Capture a variable
+    outer_var = tf.Variable(np.array([3.], dtype=np.float32))
+    x = np.array([2., 3., 4.], dtype=np.float32)
+
+    def fun_tf(x):
+      return x * tf.broadcast_to(outer_var, x.shape) + 1.
+
+    hlo = tf.function(fun_tf, jit_compile=True, autograph=False).experimental_get_compiler_ir(x)()
+    self.assertIn("(arg0.1: f32[3], arg1.2: f32[1]) -> f32[3]", hlo)
+
+    # Capture a constant
+    outer_ct = np.array([3.], dtype=np.float32)
+    x = np.array([2., 3., 4.], dtype=np.float32)
+
+    def fun_tf(x):
+      return x * tf.broadcast_to(outer_ct, x.shape) + 1.
+
+    hlo = tf.function(fun_tf, jit_compile=True, autograph=False).experimental_get_compiler_ir(x)()
+    self.assertIn("(arg0.1: f32[3]) -> f32[3]", hlo)
+
+    # Call get_compiler_ir in a function context
+    x = np.array([2., 3., 4.], dtype=np.float32)
+
+    def fun_tf_outer(x):
+      x_const = tf.constant(0, shape=x.shape, dtype=x.dtype)
+      _ = tf.function(tf.math.sin, jit_compile=True, autograph=False).experimental_get_compiler_ir(x_const)()
+
+    # TODO(b/193754660)
+    # with self.assertRaisesRegex(
+    #     TypeError, "An op outside of the function building code is being passed"):
+    #   tf.function(fun_tf_outer)(x)
+    #
+    # with self.assertRaisesRegex(
+    #     TypeError, "An op outside of the function building code is being passed"):
+    #   tf.function(fun_tf_outer, jit_compile=True)(x)
+
+    # Call get_concrete_function in a graph context
+    def fun_tf_outer_2(x):
+      _ = tf.function(tf.math.sin, jit_compile=True).get_concrete_function(tf.TensorSpec(x.shape, x.dtype))
+      return x
+
+    # Outside of a function context, this works.
+    _ = tf.function(fun_tf_outer_2)(x)
+    _ = tf.function(fun_tf_outer_2, jit_compile=True)(x)
+
+  def test_repro_193754660(self):
+    # Try to reproduce b/193754660. I can't.
+    # We have to have tf.function(jax2tf.convert(jax2tf.call_tf(f_tf))).
+    # The get_compiler_ir will indeed fail for f_tf. Then we try to use
+    # shape inference for f_tf.
+    # I thought to use a f_tf that uses an op without shape inference, e.g.,
+    # tfxla.gather. If we wash it through a saved_model I expect that shape
+    # inference would not work on it. Instead, shape inference works!!!
+    x = np.array([0, 1, 2, 3, 4, 5], dtype=np.int32)
+    def f_jax(x):
+      return x[1]
+    f_tf = jax2tf.convert(f_jax)
+    f_tf_rt, _ = tf_test_util.SaveAndLoadFunction(f_tf, input_args=[x])
+    f_jax2 = jax2tf.call_tf(f_tf_rt)
+    f_tf2 = jax2tf.convert(f_jax2)
+    res = tf.function(f_tf2, autograph=False)(x)
+    self.assertAllClose(res.numpy(), f_jax(x))
+
+  def test_effectful(self):
+    x = np.ones((3,), dtype=np.float32)
+    lower_effect = jax.jit(jax2tf.call_tf(tf.math.sin, has_side_effects=True)).lower(x)
+    self.assertNotEmpty(lower_effect._lowering.compile_args["unordered_effects"])
+
+    lower_no_effect = jax.jit(jax2tf.call_tf(tf.math.sin, has_side_effects=False)).lower(x)
+    self.assertEmpty(lower_no_effect._lowering.compile_args["unordered_effects"])
+
+  def test_module_documentation(self):
+    def cos_tf(x):
+      return tf.math.cos(x)
+
+    # Compute cos with TF and sin with JAX
+    def cos_tf_sin_jax(x):
+      return jax.numpy.sin(jax2tf.call_tf(cos_tf)(x))
+
+    # Calls `cos_tf` in TF eager mode
+    x = np.float32(1.)
+    cos_tf_sin_jax(x)
+
+    # Compiles `cos_tf` using TF and embeds the XLA computation into the JAX
+    # XLA computation (containing `sin`). The XLA compiler may even be able to
+    # fuse through JAX-TF computations.
+    jax.jit(cos_tf_sin_jax)(x)
+
+    # Uses TF gradient for `cos_tf` and JAX gradient for `sin`
+    jax.grad(cos_tf_sin_jax)(x)
+
+    logging.info(jax.make_jaxpr(cos_tf_sin_jax)(x))
+    logging.info(jax.xla_computation(cos_tf_sin_jax)(x).as_hlo_text())
+
+  def test_tf_gather(self):
+    """tf_gather gradient output is tf.IndexSlices."""
+    operand = jnp.array(np.random.uniform(size=(100, 128)))
+    indices = jnp.array(np.random.randint(low=0, high=100, size=(4000,)))
+
+    @tf.function(jit_compile=True, autograph=False)
+    def fun_tf(operand, indices):
+      return tf.experimental.numpy.std(tf.gather(operand, indices))
+
+    fun_jax = jax2tf.call_tf(fun_tf)
+    grad_fun_jax = jax.grad(fun_jax)
+    grad_res = grad_fun_jax(operand, indices)
+    self.assertEqual(grad_res.shape, (100, 128))
+
+  def test_output_shape_dtype_none(self):
+    x = jnp.zeros((10), dtype=jnp.float32)
+
+    @tf.function(jit_compile=True, autograph=False)
+    def fun_tf(x):  # pylint: disable=unused-argument
+      return
+
+    fun_jax_1 = jax2tf.call_tf(fun_tf, output_shape_dtype=None)
+    fun_jax_2 = jax2tf.call_tf(fun_tf)
+    self.assertIsNone(fun_jax_1(x))
+    self.assertIsNone(fun_jax_2(x))
+    fun_jax_3 = jax2tf.call_tf(
+        fun_tf, output_shape_dtype=jax.ShapeDtypeStruct((10,), jnp.float32)
+    )
+    with self.assertRaisesRegex(
+        ValueError,
+        "The pytree of the TensorFlow function results does not match the"
+        " pytree of the declared output_shape_dtype",
+    ):
+      _ = fun_jax_3(x)
+
+  def test_output_shape_dtype_not_none(self):
+    x = jnp.zeros((10), dtype=jnp.float32)
+
+    @tf.function(jit_compile=True, autograph=False)
+    def fun_tf(x):
+      return x
+
+    fun_jax_1 = jax2tf.call_tf(
+        fun_tf, output_shape_dtype=jax.ShapeDtypeStruct((10,), jnp.float32)
+    )
+    fun_jax_2 = jax2tf.call_tf(fun_tf)
+    self.assertAllClose(fun_jax_1(x), fun_jax_2(x))
+
+    fun_jax_3 = jax2tf.call_tf(fun_tf, output_shape_dtype=None)
+    with self.assertRaisesRegex(
+        ValueError,
+        "The pytree of the TensorFlow function results does not match the"
+        " pytree of the declared output_shape_dtype",
+    ):
+      _ = fun_jax_3(x)
+
+  def test_multi_platform(self):
+    def tf_fun(x):
+      return tf.math.sin(x)
+
+    def f_jax(x):
+      return jnp.cos(jax2tf.call_tf(tf_fun)(jnp.cos(x)))
+    x = np.arange(12, dtype=np.float32).reshape((3, 4))
+
+    # Find platforms that are available for both JAX and TF
+    # Pick one device from each available platform
+    jax_platforms = []
+    for backend in ["cpu", "gpu", "tpu"]:
+      try:
+        devices = jax.devices(backend)
+      except RuntimeError:
+        devices = []
+      if devices:
+        jax_platforms.append(devices[0].platform)
+
+    jax_and_tf_platforms = (
+      set(jax_platforms) & {d.device_type.lower()
+                            for d in self.__class__.tf_devices})
+
+    lowering_platforms = ("tpu", "cpu", "cuda")
+
+    exp = export.export(f_jax,
+                        lowering_platforms=lowering_platforms)(x)
+    for jax_platform in jax_and_tf_platforms:
+      with self.subTest(jax_platform):
+        jax_device = jax.devices(jax_platform)[0]
+        x_device = jax.device_put(x, jax_device)
+        logging.info("Running harness natively on %s", jax_device)
+        native_res = f_jax(x_device)
+        logging.info("Running exported harness on %s", jax_device)
+        exported_res = export.call(exp)(x_device)
+        self.assertAllClose(native_res, exported_res)
+
+  def test_multi_platform_call_tf_graph(self):
+    def tf_fun(x):
+      return tf.math.sin(x)
+
+    def f_jax(x):
+      return jnp.cos(jax2tf.call_tf(tf_fun,
+                                    call_tf_graph=True,
+                                    ordered=True)(jnp.cos(x)))
+    x = np.arange(12, dtype=np.float32).reshape((3, 4))
+    # When we use call_tf_graph we can serialize for multiple platforms
+    lowering_platforms = ("tpu", "cpu", "cuda")
+    # We must use jax2tf.convert to run a call_tf(call_tf_graph)
+    # TODO(necula): if we remove the tf.function and we have multiple platforms
+    # then we attempt to lower call_tf multiple times and only the first
+    # lowering will have the proper side effects for the function_list.
+    f_tf = tf.function(jax2tf.convert(
+      f_jax,
+      native_serialization=True,
+      native_serialization_platforms=lowering_platforms))
+    for tf_device in self.__class__.tf_devices:
+      with self.subTest(tf_device.device_type):
+        logging.info(
+          f"Running on tf_device = {tf_device} of device_type = {tf_device.device_type}")
+        with tf.device(tf_device):
+          res = f_tf(x)
+        self.assertAllClose(res, f_jax(x))
+
+  @parameterized.named_parameters(
+      {"testcase_name": f"_type={type_.__name__}", "type_": type_}
+      for type_ in dlpack.SUPPORTED_DTYPES
+  )
+  def test_avoid_copy_between_gpu_and_cpu(self, type_):
+    try:
+      gpu_devices = jax.devices("gpu")
+    except RuntimeError:
+      gpu_devices = []
+    if not gpu_devices:
+      raise unittest.SkipTest("Test requires a GPU device.")
+
+    def tf_fun(x):
+      if type_ == jnp.bool_:
+        return tf.math.logical_or(x, True)
+      else:
+        return x + 1
+
+    jax_array_on_gpu = jnp.zeros([1], type_, device=gpu_devices[0])
+
+    # Since the input array is already on a GPU device, we expect that no memory
+    # copy occurs between GPU and CPU. Thus, we expect no errors raised by the
+    # transfer guard.
+    # There are two exceptions:
+    # First, when dtype is "int32". This is because almost all TensorFlow
+    # kernels for GPU devices keep int32 tensors in host memory.
+    # (https://github.com/tensorflow/tensorflow/blob/4eb3e36d1b0cd511e1677e740bd093f42365cf9f/tensorflow/python/eager/pywrap_tensor.cc#L352-L354)
+    # Hence, for "int32", we do expect a "host-to-device" copy.
+    # Second, when using PJRT C API runtime. This is because it currently skips dlpack
+    # to workaround "PJRT C API does not support GetDefaultLayout" runtime error.
+    # https://github.com/openxla/xla/blob/762bde36adf22792e91c38fe87cabe5af05bfadc/xla/pjrt/pjrt_c_api_client.h#L285-L289
+    @contextlib.contextmanager
+    def _transfer_guard(guard_level):
+      with contextlib.ExitStack() as stack:
+        stack.enter_context(jax.transfer_guard_device_to_device(guard_level))
+        stack.enter_context(jax.transfer_guard_device_to_host(guard_level))
+        if type_ != jnp.int32:
+          stack.enter_context(jax.transfer_guard_host_to_device(guard_level))
+        yield
+
+    with _transfer_guard("disallow_explicit"):
+      jax2tf.call_tf(tf_fun)(jax_array_on_gpu)
+
+
+class RoundTripToJaxTest(tf_test_util.JaxToTfTestCase):
+  "Reloading output of jax2tf into JAX with call_tf"
+  def setUp(self):
+    if tf is None:
+      raise unittest.SkipTest("Test requires tensorflow")
+    # TODO(b/171320191): this line works around a missing context initialization
+    # bug in TensorFlow.
+    _ = tf.add(1, 1)
+    super().setUp()
+
+  def test_simple(self):
+    f_jax = jnp.sin
+    f_jax_rt = jax2tf.call_tf(jax2tf.convert(f_jax))
+    x = np.float32(0.7)
+    self.assertAllClose(f_jax(x), f_jax_rt(x))
+
+  def test_pytree(self):
+    def f_jax(x):  # x: dict(a=f32, b=f32)
+      return dict(a=x["a"]+1., b=x)
+    x = dict(a=0.7, b=0.8)
+    f_jax_rt = jax2tf.call_tf(jax2tf.convert(f_jax))
+    self.assertAllClose(f_jax(x), f_jax_rt(x))
+
+  def test_custom_grad(self):
+    @jax.custom_vjp
+    def f(x):
+      return x * x
+
+    # f_fwd: a -> (b, residual)
+    def f_fwd(x):
+      return f(x), np.float32(3.) * x
+    # f_bwd: (residual, CT b) -> [CT a]
+    def f_bwd(residual, ct_b):
+      return residual * ct_b,
+
+    f.defvjp(f_fwd, f_bwd)
+
+    f_rt = jax2tf.call_tf(jax2tf.convert(f, with_gradient=True))
+    x = np.float32(0.7)
+    self.assertAllClose(f(x), f_rt(x))
+    self.assertAllClose(jax.grad(f)(x), jax.grad(f_rt)(x))
+
+  def test_shape_poly(self):
+    f_jax = jnp.sin
+    f_jax_rt = jax2tf.call_tf(jax2tf.convert(f_jax,
+                                             polymorphic_shapes=["(b, ...)"]))
+    x = np.array([0.7, 0.8], dtype=np.float32)
+    self.assertAllClose(f_jax(x), f_jax_rt(x))
+
+  def test_saved_model_simple(self):
+    x = np.array([0.7, 0.8], dtype=np.float32)
+    def f_jax(x):
+      return jnp.sin(x)
+
+    f_tf = jax2tf.convert(f_jax)
+    restored_tf, _ = tf_test_util.SaveAndLoadFunction(f_tf, input_args=[x])
+    restored_jax = jax2tf.call_tf(restored_tf)
+    self.assertAllClose(f_jax(x), restored_jax(x))
+
+  def test_saved_model_variables(self):
+    param = np.array([1., 2.], dtype=np.float32)
+    x = np.array([0.7, 0.8], dtype=np.float32)
+    def f_jax(param, x):
+      return jnp.sin(x) + jnp.cos(param)
+
+    param_v = tf.Variable(param)
+    f_tf = jax2tf.convert(f_jax)
+    _, restored_model = tf_test_util.SaveAndLoadFunction(
+        lambda x: f_tf(param_v, x),
+        input_args=[x],
+        variables=[param_v])
+    restored_jax = jax2tf.call_tf(restored_model.f)
+    self.assertAllClose(f_jax(param, x), restored_jax(x))
+    self.assertAllClose(f_jax(param, x), jax.jit(restored_jax)(x))
+
+  def test_saved_model_shape_poly(self):
+    tracing_count = 0
+    x = np.array([0.7, 0.8], dtype=np.float32)
+    def f_jax(x):
+      nonlocal tracing_count
+      tracing_count += 1
+      return jnp.sin(x)
+
+    f_tf = jax2tf.convert(f_jax, polymorphic_shapes=["(b, ...)"])
+    res_jax = f_jax(x)
+    self.assertEqual(1, tracing_count)
+    # Will trace twice, it seems. Once to get the result signature, and once again
+    # for the actual saving.
+    restored_f, _ = tf_test_util.SaveAndLoadFunction(
+        f_tf, input_signature=[tf.TensorSpec([None], x.dtype)])
+    self.assertGreaterEqual(tracing_count, 2)
+    tracing_count = 0
+    f_jax_rt = jax2tf.call_tf(restored_f)
+    self.assertAllClose(res_jax, f_jax_rt(x))
+    # Ensure that restored_f works at other batch size as well
+    y = np.concatenate([x, x])
+    self.assertEqual(0, tracing_count)
+    res_jax_y = f_jax(y)
+    self.assertEqual(1, tracing_count)
+    # No more tracing for f_jax_rt
+    self.assertAllClose(res_jax_y, f_jax_rt(y))
+    self.assertEqual(1, tracing_count)
+
+  def test_custom_grad_saved_model(self):
+
+    @jax.custom_vjp
+    def f(x):
+      return x * x
+
+    # f_fwd: a -> (b, residual)
+    def f_fwd(x):
+      return f(x), np.float32(3.) * x
+    # f_bwd: (residual, CT b) -> [CT a]
+    def f_bwd(residual, ct_b):
+      return residual * ct_b,
+
+    f.defvjp(f_fwd, f_bwd)
+    def g(x):
+      return jnp.sum(f(x))
+
+    g_tf, _ = tf_test_util.SaveAndLoadFunction(
+        jax2tf.convert(g, with_gradient=True),
+        input_signature=[tf.TensorSpec(shape=(1,), dtype=tf.float32)],
+    )
+    g_rt = jax2tf.call_tf(g_tf)
+    x = np.array([0.7], dtype=np.float32)
+    self.assertAllClose(g(x), g_rt(x))
+    self.assertAllClose(jax.grad(g)(x), jax.grad(g_rt)(x))
+
+  def test_without_gradient_saved_model(self):
+    # Explicitly with_gradient=False
+    f_jax = jnp.sum
+
+    x = np.array([0.7, 0.8], dtype=np.float32)
+    f_tf, _ = tf_test_util.SaveAndLoadFunction(
+        jax2tf.convert(f_jax, with_gradient=False),
+        input_args=[x])
+    f_rt = jax2tf.call_tf(f_tf)
+
+    self.assertAllClose(f_jax(x), f_rt(x))
+    with self.assertRaisesRegex(Exception,
+                                "Gradient explicitly disabled.*jax2tf-converted function does not support gradients. Use `with_gradient` parameter to enable gradients"):
+      jax.grad(f_rt)(x)
+
+  def test_saved_model_no_gradients(self):
+    # Save without gradients
+    f_jax = jnp.sum
+
+    x = np.array([0.7, 0.8], dtype=np.float32)
+    f_tf, _ = tf_test_util.SaveAndLoadFunction(
+        jax2tf.convert(f_jax, with_gradient=True), input_args=[x],
+        save_gradients=False)
+    f_rt = jax2tf.call_tf(f_tf)
+
+    self.assertAllClose(f_jax(x), f_rt(x))
+    # TODO: clean this up b/191117111: it should fail with a clear error
+    # The following results in a confusing error:
+    # TypeError: tf.Graph captured an external symbolic tensor.
+    with self.assertRaises(TypeError):
+      _ = jax.grad(f_rt)(x)
+
+  def test_call_tf_under_function_context(self):
+    def fun_jax(x, y):
+      z = jax2tf.call_tf(tf.math.sin)(x) + jnp.cos(y)
+      return z
+
+    x = np.array([-1.0, 0.0, 1.0], dtype=np.float32)
+    y = np.array([-0.5, 0.0, 0.5], dtype=np.float32)
+
+    converted_fun = tf.function(
+        jax2tf.convert(fun_jax, native_serialization=True)
+    )
+    expected = np.sin(x) + np.cos(y)
+    res = tf.function(converted_fun, jit_compile=True, autograph=False)(x, y)
+    self.assertAllClose(expected, res.numpy(), atol=1e-5, rtol=1e-5)
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name=f"_{dtype.__name__}",
+          dtype=dtype,
+      )
+      for dtype in set(jtu.dtypes.all_floating)
+  )
+  def test_all_floating_input_gradient(self, dtype):
+    def tf_f(x):
+      res = tf.math.sin(x)
+      return tf.reduce_sum(res)
+
+    jax_f = jax2tf.call_tf(tf_f)
+    tf_f_rt = jax2tf.convert(jax_f)
+    x = jnp.array([5.0, 6.0, 7.0]).astype(dtype)
+
+    def assert_all_close_support_bfloat16(baseline, candidate):
+      def conversion(x):
+        # convert scalar to array and bfloat16 to float32
+        # to support self.assertAllClose numpy array comparison.
+        if x.shape == tf.TensorShape([]):
+          x = tf.convert_to_tensor([x])
+        if dtype == jnp.float16:
+          x = tf.cast(x, tf.float32)
+        return x
+
+      baseline = jax.tree_util.tree_map(conversion, baseline)
+      candidate = jax.tree_util.tree_map(conversion, candidate)
+      tol = (
+          1e-2
+          if jtu.test_device_matches(["tpu"]) and dtype == np.float16
+          else None
+      )
+      self.assertAllClose(baseline, candidate, atol=tol, rtol=tol)
+
+    # Eager mode
+    assert_all_close_support_bfloat16(tf_f(x), tf_f_rt(x))
+
+    # Compiled function mode
+    assert_all_close_support_bfloat16(
+        tf.function(tf_f)(x), tf.function(tf_f_rt)(x)
+    )
+
+    # Compiled function mode with jit_compiled=True
+    assert_all_close_support_bfloat16(
+        tf.function(tf_f, jit_compile=True)(x),
+        tf.function(tf_f_rt, jit_compile=True)(x),
+    )
+
+    # RoundTrip test for the gradient
+    grad_fun_jax = jax.grad(jax2tf.call_tf(tf_f))
+    grad_fun_jax_rt = jax2tf.call_tf(jax2tf.convert(grad_fun_jax))
+
+    # Eager mode
+    assert_all_close_support_bfloat16(grad_fun_jax(x), grad_fun_jax_rt(x))
+
+    # Jit mode
+    assert_all_close_support_bfloat16(
+        jax.jit(grad_fun_jax)(x), jax.jit(grad_fun_jax_rt)(x)
+    )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name=f"_{dtype.__name__}",
+          dtype=dtype,
+      )
+      for dtype in set(jtu.dtypes.complex)
+  )
+  def test_complex_input_gradient(self, dtype):
+    def tf_f(x):
+      res = tf.math.sin(x)
+      return tf.reduce_sum(res)
+
+    x = jnp.array([(5.0 + 4.0j), (6.0 + 3.0j), (7.0 + 8.0j)]).astype(dtype)
+
+    jax_f = jax2tf.call_tf(tf_f)
+    tf_f_rt = jax2tf.convert(jax_f)
+
+    # Eager mode
+    self.assertAllClose(tf_f(x), tf_f_rt(x))
+
+    # tf.function context
+    self.assertAllClose(tf.function(tf_f)(x), tf.function(tf_f_rt)(x))
+
+    # tf.function context with jit_compiled=True
+    self.assertAllClose(
+        tf.function(tf_f, jit_compile=True)(x),
+        tf.function(tf_f_rt, jit_compile=True)(x),
+    )
+
+    # RoundTrip test for the gradient
+    grad_fun_jax = jax.grad(jax2tf.call_tf(tf_f), holomorphic=True)
+    grad_fun_jax_rt = jax2tf.call_tf(jax2tf.convert(grad_fun_jax))
+
+    # Eager mode
+    self.assertAllClose(grad_fun_jax(x), grad_fun_jax_rt(x))
+
+    # Jit mode
+    self.assertAllClose(jax.jit(grad_fun_jax)(x), jax.jit(grad_fun_jax_rt)(x))
+
+
+class RoundTripToTfTest(tf_test_util.JaxToTfTestCase):
+  "Reloading output of call_tf into TF with jax2tf."
+
+  def setUp(self):
+    if tf is None:
+      raise unittest.SkipTest("Test requires tensorflow")
+    # TODO(b/171320191): this line works around a missing context initialization
+    # bug in TensorFlow.
+    _ = tf.add(1, 1)
+    super().setUp()
+
+  def test_alternate(self):
+    # Alternate sin/cos with sin in TF and cos in JAX
+    f_tf_inner = tf.math.sin
+    def f_jax(x_jax):
+      y_jax = jnp.cos(x_jax)
+      z_jax = jax2tf.call_tf(f_tf_inner)(y_jax)
+      return jnp.cos(z_jax)
+    def f_tf_outer(x_tf):
+      y_tf = tf.math.sin(x_tf)
+      z_tf = jax2tf.convert(f_jax)(y_tf)
+      return tf.math.sin(z_tf)
+
+    x = np.float32(0.7)
+
+    self.assertAllClose(np.sin(np.cos(np.sin(np.cos(np.sin(x))))),
+                        f_tf_outer(x).numpy())
+    xv = tf.Variable(x)
+    with tf.GradientTape() as tape:
+      res = f_tf_outer(xv)
+    g_tf = tape.gradient(res, xv)
+    _, gf = tf_test_util.ComputeTfValueAndGrad(f_tf_outer, (x,))
+    # Eager
+    expected_res = np.sin(np.cos(np.sin(np.cos(np.sin(x)))))
+    self.assertAllClose(expected_res, f_tf_outer(x).numpy())
+
+    # Gradient
+    expected_grad = (np.cos(np.cos(np.sin(np.cos(np.sin(x))))) *
+                     np.sin(np.sin(np.cos(np.sin(x)))) *
+                     np.cos(np.cos(np.sin(x))) *
+                     np.sin(np.sin(x)) *
+                     np.cos(x))
+    self.assertAllClose(expected_grad, g_tf.numpy())
+
+    # Graph
+    self.assertAllClose(expected_res,
+                        tf.function(f_tf_outer, autograph=False)(x).numpy())
+
+    # Compiled
+    self.assertAllClose(expected_res,
+                        tf.function(f_tf_outer, autograph=False,
+                                    jit_compile=True)(x).numpy())
+
+  def test_saved_model(self):
+    x = np.array([.7, .8], dtype=np.float32)
+    def fun_tf(x):
+      return tf.math.sin(x)
+    def fun_jax(x):
+      return jax2tf.call_tf(fun_tf)(x)
+
+    # Now convert and save to SavedModel
+    fun_tf_rt = jax2tf.convert(fun_jax)
+    res = fun_tf_rt(x)
+    self.assertAllClose(np.sin(x), res.numpy())
+
+    res = tf.function(fun_tf_rt, autograph=False)(x)
+    self.assertAllClose(np.sin(x), res.numpy())
+
+    res = tf.function(fun_tf_rt, jit_compile=True, autograph=False)(x)
+    self.assertAllClose(np.sin(x), res.numpy())
+
+    reloaded_f, _ = tf_test_util.SaveAndLoadFunction(
+        fun_tf_rt, input_args=[x])
+    res = reloaded_f(x)
+    self.assertAllClose(np.sin(x), res.numpy())
+
+  def test_saved_model_polymorphic_input_static_output(self):
+    x = np.array([.7, .8], dtype=np.float32)
+    def fun_tf(x):
+      return tf.math.reduce_sum(tf.math.sin(x))
+    def fun_jax(x):
+      return jax2tf.call_tf(fun_tf)(x)
+
+    # Now convert and save to SavedModel
+    fun_tf_rt = jax2tf.convert(fun_jax)
+    res = fun_tf_rt(x)
+    self.assertAllClose(fun_tf(x), res.numpy())
+
+    res = tf.function(fun_tf_rt, autograph=False)(x)
+    self.assertAllClose(fun_tf(x), res.numpy())
+
+    res = tf.function(fun_tf_rt, jit_compile=True, autograph=False)(x)
+    self.assertAllClose(fun_tf(x), res.numpy())
+
+    reloaded_f, _ = tf_test_util.SaveAndLoadFunction(
+        fun_tf_rt, input_args=[x])
+    res = reloaded_f(x)
+    self.assertAllClose(fun_tf(x), res.numpy())
+
+  def test_function_dynamic_shape(self):
+    # Call a function for which shape inference does not give an output
+    # shape.
+    x = np.array([-1, 0, 1], dtype=np.int32)
+    def fun_tf(x):  # x:i32[3]
+      # The shape depends on the value of x
+      return tf.cond(x[0] >= 0, lambda: x, lambda: x[1:])
+
+    # Call in eager mode. Should work!
+    res1 = jax2tf.call_tf(fun_tf)(x)
+    expected = x[1:]
+    self.assertAllClose(expected, res1, check_dtypes=False)
+
+    # Now under jit, should fail because the function is not compilable
+    with self.assertRaisesRegex(ValueError, _call_tf_dynamic_shape_error):
+      fun_jax = jax.jit(jax2tf.call_tf(fun_tf))
+      fun_jax(x)
+
+    # TODO(necula): this should work in op-by-op mode, but it fails because
+    # jax2tf.convert does abstract evaluation.
+    with self.assertRaisesRegex(ValueError, _call_tf_dynamic_shape_error):
+      fun_tf_rt = jax2tf.convert(jax2tf.call_tf(fun_tf))
+      fun_tf_rt(x)
+
+  @_parameterized_jit
+  def test_shape_poly_static_output_shape(self, with_jit=True):
+    if jax.config.jax2tf_default_native_serialization:
+      raise unittest.SkipTest("TODO(b/268386622): call_tf with shape polymorphism and native serialization.")
+    x = np.array([0.7, 0.8], dtype=np.float32)
+
+    def fun_tf(x):
+      return tf.math.reduce_sum(tf.math.sin(x))
+
+    fun_jax = jax2tf.call_tf(fun_tf)
+    fun_tf_rt = _maybe_tf_jit(with_jit,
+        jax2tf.convert(fun_jax, polymorphic_shapes=["b, ..."]))
+    self.assertAllClose(fun_tf(x), fun_tf_rt(x))
+
+  @_parameterized_jit
+  def test_shape_poly(self, with_jit=False):
+    if jax.config.jax2tf_default_native_serialization:
+      raise unittest.SkipTest("TODO(b/268386622): call_tf with shape polymorphism and native serialization.")
+    x = np.array([7, 8, 9, 10], dtype=np.float32)
+    def fun_jax(x):
+      y = jax2tf.call_tf(tf.math.sin,
+                         output_shape_dtype=jax.ShapeDtypeStruct(x.shape, x.dtype))(x)
+      z = jnp.cos(y)
+      w = jax2tf.call_tf(lambda z: tf.concat([z, z], axis=0),
+                         output_shape_dtype=jax.ShapeDtypeStruct((2 * z.shape[0],), z.dtype))(z)
+      assert w.shape[0] == 2 * x.shape[0]
+      return w
+
+    fun_tf_rt = _maybe_tf_jit(with_jit,
+        jax2tf.convert(fun_jax, polymorphic_shapes=["b, ..."]))
+    res_tf = fun_tf_rt(x)
+    self.assertAllClose(fun_jax(x), res_tf)
+
+  @_parameterized_jit
+  def test_shape_poly_pytree_result(self, with_jit=True):
+    if jax.config.jax2tf_default_native_serialization:
+      raise unittest.SkipTest("TODO(b/268386622): call_tf with shape polymorphism and native serialization.")
+    x = np.array([7, 8, 9, 10], dtype=np.float32)
+    def fun_jax(x):
+      # Returns a tuple
+      y = jax2tf.call_tf(lambda x: (x, tf.concat([x, x], axis=0)),
+          output_shape_dtype=(jax.ShapeDtypeStruct(x.shape, x.dtype),
+                              jax.ShapeDtypeStruct((2 * x.shape[0],), x.dtype)))(x)
+      assert y[0].shape[0] == x.shape[0]
+      assert y[1].shape[0] == 2 * x.shape[0]
+      return y
+
+    fun_tf_rt = _maybe_tf_jit(with_jit,
+        jax2tf.convert(fun_jax, polymorphic_shapes=["b, ..."]))
+    res_tf = fun_tf_rt(x)
+    self.assertAllClose(fun_jax(x), res_tf)
+
+  @_parameterized_jit
+  def test_shape_poly_error_no_output_shape_dtype(self, with_jit=True):
+    x = np.array([7, 8, 9, 10], dtype=np.float32)
+    def fun_jax(x):
+      return jax2tf.call_tf(tf.math.sin)(x)
+
+    fun_tf_rt = _maybe_tf_jit(with_jit,
+        jax2tf.convert(fun_jax, polymorphic_shapes=["b, ..."]))
+    with self.assertRaisesRegex(ValueError, _call_tf_dynamic_shape_error):
+      fun_tf_rt(x)
+
+  @_parameterized_jit
+  def test_shape_poly_error_mismatch_output_shape_dtype_tree(self, with_jit=False):
+    x = np.array([7, 8, 9, 10], dtype=np.float32)
+    def fun_jax(x):
+      return jax2tf.call_tf(tf.math.sin,
+          output_shape_dtype=(jax.ShapeDtypeStruct(x.shape, x.dtype),
+                              jax.ShapeDtypeStruct(x.shape, x.dtype)))(x)
+
+    fun_tf_rt = _maybe_tf_jit(with_jit,
+        jax2tf.convert(fun_jax, polymorphic_shapes=["b, ..."]))
+
+    with self.assertRaisesRegex(
+        ValueError,
+        "The pytree of the TensorFlow function results does not match the pytree of the declared output_shape_dtype"):
+      fun_tf_rt(x)
+
+  @parameterized.named_parameters(
+      _named_test(with_jit=with_jit, kind=kind)
+      for with_jit in [True, False]
+      for kind in ["bad_rank", "bad_dim", "bad_dtype", "bad_dtype_x64"])
+  def test_shape_poly_error_mismatch_output_shape_dtype(self, with_jit=False, kind="bad_rank"):
+    x = np.array([7, 8, 9, 10], dtype=np.float32)
+
+    if kind == "bad_rank":
+      def fun_jax(x):
+        return jax2tf.call_tf(lambda x: x,
+                              # Wrong shape rank
+                              output_shape_dtype=jax.ShapeDtypeStruct((), x.dtype))(x)
+    elif kind == "bad_dim":
+      def fun_jax(x):
+        bad_shape = (5 + x.shape[0],)
+        y = jax2tf.call_tf(lambda x: x,
+                           # Wrong dimension
+                           output_shape_dtype=jax.ShapeDtypeStruct(bad_shape, x.dtype))(x)
+        # JAX will believe that the following is Ok, leading to downstream error in TF
+        return y + jnp.ones(bad_shape, dtype=x.dtype)
+    elif kind == "bad_dtype":
+      def fun_jax(x):
+        return jax2tf.call_tf(lambda x: x,
+                              output_shape_dtype=jax.ShapeDtypeStruct(x.shape, np.int32))(x)
+    elif kind == "bad_dtype_x64":
+      def fun_jax(x):
+        return jax2tf.call_tf(lambda x: x * np.float64(3.),
+                              output_shape_dtype=jax.ShapeDtypeStruct(x.shape, np.float64))(x)
+    else:
+      assert False
+    expect_ex = ValueError
+    expect_error = r"The shapes or dtypes returned by the TensorFlow function do not match the declared output_shape_dtype"
+
+    # Call without shape polymorphism
+    fun_tf_rt = _maybe_tf_jit(with_jit, jax2tf.convert(fun_jax))
+    with self.assertRaisesRegex(expect_ex, expect_error):
+      fun_tf_rt(x)
+
+    # Now with shape polymorphism
+    if kind == "bad_dim" and with_jit:
+      # TODO: in jit more the error pops up later, at AddV2
+      expect_error = "Dimensions must be equal, but are 4 and 9 for .* AddV2"
+    if kind == "bad_dim" and jax.config.jax2tf_default_native_serialization:
+      # TODO(b/268386622): call_tf with shape polymorphism and native serialization.
+      expect_error = "Error compiling TensorFlow function"
+    fun_tf_rt = _maybe_tf_jit(with_jit,
+        jax2tf.convert(fun_jax, polymorphic_shapes=["b, ..."]))
+    with self.assertRaisesRegex(expect_ex, expect_error):
+      fun_tf_rt(x)
+
+  def test_inner_native_serialization(self):
+    # Two nested jax2tf, the inner one being with native serialization
+    x = np.ones((3,), dtype=np.float32)
+    def f_inner_jax(x):
+      return jnp.sin(x)
+    def f_outer_jax(x):
+      f_inner_tf = jax2tf.convert(f_inner_jax, native_serialization=True)
+      return jnp.cos(jax2tf.call_tf(f_inner_tf)(x))
+
+    f_outer_tf = tf.function(
+        jax2tf.convert(f_outer_jax, native_serialization=False),
+        autograph=False)
+    f_outer_graph = str(f_outer_tf.get_concrete_function(tf.convert_to_tensor(x)).graph.as_graph_def())
+    # Quick way to check that there is an XlaCallModule op, and a Cos op, but no Sin op
+    self.assertIn('op: "Cos"', f_outer_graph)
+    self.assertIn('op: "XlaCallModule"', f_outer_graph)
+    self.assertNotIn('op: "Sin"', f_outer_graph)
+
+  @parameterized.named_parameters(
+      _named_test(f2_function=f2_function, f2_saved_model=f2_saved_model,
+                  f4_function=f4_function, f4_saved_model=f4_saved_model)
+      for f2_function in [True, False]
+      for f2_saved_model in [True, False]
+      for f4_function in [True, False]
+      for f4_saved_model in [True, False])
+  def test_several_round_trips(self,
+                               f2_function=False, f2_saved_model=False,
+                               f4_function=False, f4_saved_model=False):
+    if (f2_saved_model and
+        f4_saved_model and
+        not jax.config.jax2tf_default_native_serialization):
+      # TODO: Getting error Found invalid capture Tensor("jax2tf_vjp/jax2tf_arg_0:0", shape=(), dtype=float32) when saving custom gradients
+      # when saving f4, but only with non-native serialization.
+      raise unittest.SkipTest("TODO: error invalid capture when saving custom gradients")
+    x = np.array(.7, dtype=np.float32)
+    # f(n)(x) = 2. * x^n
+    def f(n):
+      def fn(x):
+        acc = np.array(2., dtype=x.dtype)
+        for i in range(n):
+          acc *= x
+        return acc
+      return fn
+
+    f2_tf = lambda x: x * jax2tf.convert(f(1))(x)
+    if f2_function:
+      f2_tf = tf.function(f2_tf, autograph=False)
+    if f2_saved_model:
+      f2_tf, _ = tf_test_util.SaveAndLoadFunction(f2_tf, input_args=[x])
+
+    self.assertAllClose(f(2)(x), f2_tf(x).numpy())
+    _, (g_f2_ft,) = tf_test_util.ComputeTfValueAndGrad(f2_tf, [x])
+    self.assertAllClose(jax.grad(f(2))(x), g_f2_ft.numpy())
+
+    f3_jax = lambda x: x * jax2tf.call_tf(f2_tf)(x)
+    self.assertAllClose(f(3)(x), f3_jax(x))
+    self.assertAllClose(f(3)(x), jax.jit(f3_jax)(x))
+    self.assertAllClose(jax.grad(f(3))(x), jax.grad(f3_jax)(x))
+
+    f4_tf = lambda x: x * jax2tf.convert(f3_jax)(x)
+    self.assertAllClose(f(4)(x), f4_tf(x).numpy())
+    _, (g_f4_ft,) = tf_test_util.ComputeTfValueAndGrad(f4_tf, [x])
+    self.assertAllClose(jax.grad(f(4))(x), g_f4_ft.numpy())
+
+    if f4_function:
+      f4_tf = tf.function(f4_tf, autograph=False)
+    if f4_saved_model:
+      f4_tf, _ = tf_test_util.SaveAndLoadFunction(f4_tf, input_args=[x])
+    self.assertAllClose(f(4)(x), f4_tf(x).numpy())
+    _, (g_f4_ft,) = tf_test_util.ComputeTfValueAndGrad(f4_tf, [x])
+    self.assertAllClose(jax.grad(f(4))(x), g_f4_ft.numpy())
+
+  @classmethod
+  def _walk_stablehlo_operations(cls, op, cb):
+    """walk the stablehlo operation recursive with callback function."""
+    cb(op)
+    for region in op.operation.regions:
+      for block in region:
+        for op in block:
+          cls._walk_stablehlo_operations(op, cb)
+
+  def test_call_tf_graph(self):
+    const = tf.Variable(0.0, dtype=tf.float32)
+
+    @tf.function(jit_compile=True)
+    def tf_func_1(x):
+      return x * x + const
+
+    @tf.function
+    def tf_func_2(x, y):
+      return tf_func_1(x) + y
+
+    @tf.function
+    def tf_func_3(x, y, z):
+      return tf_func_2(x, y) + z, z
+
+    x = jnp.array(3.0, dtype=jnp.float32)
+    y = jnp.array(3.0, dtype=jnp.float32)
+    z = jnp.array(5.0, dtype=jnp.float32)
+    f_jax = jax.jit(jax2tf.call_tf(tf_func_3, call_tf_graph=False))
+    stablehlo_module = f_jax.lower(x, y, z).compiler_ir("stablehlo")
+    self.assertNotIn("stablehlo.custom_call", str(stablehlo_module))
+
+    f_jax = jax.jit(
+        jax2tf.call_tf(
+            tf_func_3,
+            call_tf_graph=True,
+        )
+    )
+    with self.assertRaisesRegex(
+        ValueError,
+        "call_tf_graph=True only support exporting by jax2tf.convert currently",
+    ):
+      stablehlo_module = f_jax.lower(x, y, z).compiler_ir("stablehlo")
+      self.assertIn("stablehlo.custom_call", str(stablehlo_module))
+
+      called_index_list = []
+
+      def _extract_info(op):
+        if op.operation.name != "stablehlo.custom_call":
+          return
+        tf_backend_config = ir.DictAttr(op.attributes["tf.backend_config"])
+        called_index = ir.IntegerAttr(tf_backend_config["called_index"]).value
+        called_index_list.append(called_index)
+
+      self._walk_stablehlo_operations(stablehlo_module, _extract_info)
+      self.assertLen(called_index_list, 1)
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="multiple_outputs",
+          tf_f=lambda x: tf.py_function(np.sin, [x], tf.float32),
+          output_shape_dtype=jax.ShapeDtypeStruct((10,), jnp.float32),
+      ),
+      dict(
+          testcase_name="zero_outputs",
+          tf_f=lambda x: print(tf.strings.length(tf.constant("hello, world"))),
+          output_shape_dtype=None,
+      ),
+  )
+  def test_call_tf_graph_non_compilable(self, tf_f, output_shape_dtype):
+    inputs = jnp.ones([10], dtype=jnp.float32)
+    called_index_list = []
+    xla_call_module_list = []
+
+    def _extract_info(op):
+      if op.operation.name != "stablehlo.custom_call":
+        return
+      tf_backend_config = ir.DictAttr(op.attributes["tf.backend_config"])
+      called_index = ir.IntegerAttr(tf_backend_config["called_index"]).value
+      called_index_list.append(called_index)
+
+    jax_f = jax2tf.call_tf(
+        tf_f,
+        call_tf_graph=True,
+        output_shape_dtype=output_shape_dtype,
+    )
+
+    # Eager mode
+    self.assertAllClose(tf_f(inputs), jax_f(inputs))
+
+    # Jit mode
+    stablehlo_module = None
+    with self.assertRaisesRegex(
+        ValueError,
+        "call_tf_graph=True only support exporting by jax2tf.convert currently",
+    ):
+      stablehlo_module = jax.jit(jax_f).lower(inputs).compiler_ir("stablehlo")
+    if stablehlo_module:
+      self.assertIn(
+          "stablehlo.custom_call @tf.call_tf_function",
+          str(stablehlo_module),
+      )
+      self.assertIn("tf.backend_config", str(stablehlo_module))
+      self._walk_stablehlo_operations(stablehlo_module, _extract_info)
+      self.assertLen(called_index_list, 1)
+
+    # Test model exporting and reloading.
+    # There is no runtime support yet so it can not run.
+    tf_f_rt = jax2tf.convert(
+        jax_f,
+        native_serialization=True,
+        with_gradient=False,
+    )
+    _, restored_model = tf_test_util.SaveAndLoadFunction(
+        tf_f_rt, input_args=[inputs]
+    )
+    func_def = restored_model.f.concrete_functions[0]
+
+    for node_def in func_def.graph.as_graph_def().node:
+      if node_def.op == "XlaCallModule":
+        xla_call_module_list.append(node_def)
+    # There is only one xla_call_module in the saved model.
+    self.assertLen(xla_call_module_list, 1)
+
+    # Check the xla_call_module version and function_list attributes.
+    xla_call_module = xla_call_module_list[0]
+    self.assertGreaterEqual(xla_call_module.attr["version"].i, 5)
+    self.assertIn("function_list", str(xla_call_module.attr))
+    xla_call_module_list.clear()
+    called_index_list.clear()
+
+    # If JAX calls same tensorflow function by `jax2tf.call_tf` twice,
+    # it should return two different tf concrete functions.
+    def jax_f_2(x):
+      res1 = jax2tf.call_tf(
+          tf_f,
+          call_tf_graph=True,
+          output_shape_dtype=output_shape_dtype,
+      )(x)
+      res2 = jax2tf.call_tf(
+          tf_f,
+          call_tf_graph=True,
+          output_shape_dtype=output_shape_dtype,
+      )(x)
+      return res1, res2
+    stablehlo_module = None
+    with self.assertRaisesRegex(ValueError, "call_tf_graph=True only support exporting by jax2tf.convert currently"):
+      stablehlo_module = jax.jit(jax_f_2).lower(inputs).compiler_ir("stablehlo")
+    if stablehlo_module:
+      self._walk_stablehlo_operations(stablehlo_module, _extract_info)
+    xla_call_module_list.clear()
+
+  def test_b279454591(self):
+    """Test case when tensorflow function returns `StatefulPartitionedCall` op."""
+    inputs = jnp.ones([10], dtype=jnp.float32)
+
+    # With one or more outputs, it is okay.
+    def tf_f(x):
+      y = tf.math.sin(3.0)
+      tf.print(y)
+      return x
+
+    jax_f = jax2tf.call_tf(
+        tf.function(tf_f),
+        call_tf_graph=True,
+    )
+    tf_f_rt = jax2tf.convert(
+        jax_f,
+        native_serialization=True,
+        with_gradient=False,
+    )
+    _, _ = tf_test_util.SaveAndLoadFunction(tf_f_rt, input_args=[inputs])
+
+    # With zero output, it return `StatefulPartitionedCall` op instead.
+    def tf_f_2():
+      y = tf.math.sin(3.0)
+      tf.print(y)
+      return
+
+    jax_f_2 = jax2tf.call_tf(tf.function(tf_f_2), call_tf_graph=True)
+    tf_f_rt_2 = jax2tf.convert(
+        jax_f_2,
+        native_serialization=True,
+        with_gradient=False,
+    )
+    _, _ = tf_test_util.SaveAndLoadFunction(tf_f_rt_2, input_args=[])
+
+  @jtu.parameterized_filterable(
+    kwargs=[dict(version=version) for version in [9]]
+  )
+  def test_call_tf_graph_ordered(self, *, version: int):
+    with config.jax_serialization_version(version):
+      logging.info(
+        "Using JAX serialization version %s",
+        jax.config.jax_serialization_version)
+
+      @tf.function
+      def tf_print(x):
+        tf.print(x)
+
+      call_tf_print = jax2tf.call_tf(
+          tf_print,
+          call_tf_graph=True,
+          ordered=True,
+      )
+
+      x = jnp.array(1.0, dtype=jnp.float32)
+
+      def body(i, x):
+        call_tf_print(x)
+        return x + 1
+
+      @jax.jit
+      def f_jax(x):
+        return jax.lax.fori_loop(0, 4, body, x)
+
+      num_custom_calls = 0
+
+      def _check_mlir_ops(op):
+        nonlocal num_custom_calls
+
+        if (
+            op.operation.name == "stablehlo.custom_call"
+            and ir.StringAttr(op.attributes["call_target_name"]).value
+            == "tf.call_tf_function"
+        ):
+          num_custom_calls += 1
+
+          # The custom call op must have `has_token_input_output` attribute.
+          tf_backend_config = ir.DictAttr(op.attributes["tf.backend_config"])
+          self.assertTrue(
+              ir.BoolAttr(tf_backend_config["has_token_input_output"]).value
+          )
+
+          # Verify that the first argument/result of the custom call op is a token
+          # type. This is a calling convention defined by `has_token_input_output`.
+          self.assertTrue(hlo.TokenType.isinstance(op.operands[0].type))
+          self.assertTrue(hlo.TokenType.isinstance(op.results[0].type))
+
+      stablehlo_module = None
+      with self.assertRaisesRegex(
+          ValueError,
+          "call_tf_graph=True only support exporting by jax2tf.convert currently",
+      ):
+        lower = f_jax.lower(x)
+        self.assertNotEmpty(lower._lowering.compile_args["ordered_effects"])
+        stablehlo_module = lower.compiler_ir("stablehlo")
+      if stablehlo_module:
+        self._walk_stablehlo_operations(stablehlo_module, _check_mlir_ops)
+        self.assertEqual(num_custom_calls, 1)
+
+      f_tf = jax2tf.convert(
+          f_jax,
+          native_serialization=True,
+          with_gradient=False,
+      )
+      _, restored_model = tf_test_util.SaveAndLoadFunction(f_tf, input_args=[x])
+
+  @jtu.parameterized_filterable(
+    kwargs=[dict(poly=poly, version=version)
+            for poly in [True, False]
+            for version in [9]]
+  )
+  def test_call_tf_ordered_dead_inputs(self, *, poly: bool, version: int):
+    with config.jax_serialization_version(version):
+      logging.info(
+        "Using JAX serialization version %s",
+        jax.config.jax_serialization_version)
+      def f_jax(x1, x_dead, x3):
+        return (x1, jax2tf.call_tf(lambda x: tf.math.sin(x), ordered=True,
+                                  call_tf_graph=True)(x3))
+      if poly:
+        polymorphic_shapes = ["b", None, None]
+      else:
+        polymorphic_shapes = None
+      f_tf = jax2tf.convert(f_jax, polymorphic_shapes=polymorphic_shapes)
+      x1 = np.arange(3, dtype=np.float32)
+      x_dead = np.arange(4, dtype=np.float32)
+      x3 = np.arange(5, dtype=np.float32)
+      self.assertAllClose(f_jax(x1, x_dead, x3),
+                          f_tf(x1, x_dead, x3))
+
+  @jtu.parameterized_filterable(
+    kwargs=[dict(ordered=ordered, version=version)
+      for ordered in [True, False]
+      for version in [9]
+    ]
+  )
+  def test_call_tf_graph_polymorphic(self, ordered: bool, version: int):
+    with config.jax_serialization_version(version):
+      logging.info(
+        "Using JAX serialization version %s",
+        jax.config.jax_serialization_version)
+
+      @tf.function(jit_compile=True, autograph=False)
+      @partial(jax2tf.convert,
+        with_gradient=False,
+        native_serialization=True,
+        polymorphic_shapes=["(b)"])
+      @jax.jit
+      def tf_f_2(x):
+        tf_f = lambda x: print(tf.strings.length(tf.constant("hello, world")))
+        jax2tf.call_tf(tf_f,
+                      call_tf_graph=True,
+                      ordered=ordered)(x)
+        return x
+
+      x = np.arange(3, dtype=np.int32)
+      _ = tf.function(tf_f_2, autograph=False).get_concrete_function(x)
+
+  # TODO(b/293927250): call_tf_graph=True only accept concrete_function. The
+  # workaround here is to set `module.call=concrete_fn.`.
+  @unittest.skip(
+      "The root cause here is because the XLACallModule.function_list attribute"
+      " depends on JAX call_tf lowering. The 2nd time tf.SavedModel TF tracing"
+      " will not trigger call_tf tracing since it was already cached. The"
+      " solution is to create the `CallTFContext` to make TF tracing and JAX"
+      " tracing work together correctly."
+  )
+  def test_call_tf_graph_save_and_load(self):
+    def jax_func(x):
+      def func_tf(x):
+        return tf.math.sin(x)
+
+      return jnp.cos(
+          jax2tf.call_tf(func_tf, output_shape_dtype=x, call_tf_graph=True)(x)
+      )
+    data_inputs = (np.array([0.5, 0.7], dtype=np.float32),)
+
+    def tf_func(the_input):
+      res = jax2tf.convert(jax_func, native_serialization=True)(the_input)
+      return tf.identity(res, name="the_result")
+
+    jit_tf_func = tf.function(
+        tf_func,
+        autograph=False,
+        jit_compile=True,
+    )
+    # The next line is necessary to reproduce this issue. It trigger TF
+    # ConcreteFunction tracing. Otherwise, you will fail with another error
+    # `Found zero restored functions for caller function`.
+    _ = jit_tf_func.get_concrete_function(*data_inputs)
+    module = tf.Module()
+    module.call = jit_tf_func # Switching to concrete_function works.
+    root_dir = self.create_tempdir()
+    saved_model_dir = os.path.join(root_dir, "saved_model")
+    tf.saved_model.save(
+        module,
+        saved_model_dir,
+        options=tf.saved_model.SaveOptions(experimental_custom_gradients=False),
+    )
+    loaded_model = tf.saved_model.load(saved_model_dir)
+    res = loaded_model.call(*data_inputs)
+    self.assertAllClose(jax_func(*data_inputs), res)
+
+
+if __name__ == "__main__":
+  absltest.main(testLoader=jtu.JaxTestLoader())