Remove explicit type stub for jax/lib/__init__.py, now that the pytype bug is fixed.
PiperOrigin-RevId: 408626593
Change-Id: I1359c30bba8552abd2eeff6f4c1fa22944e9342b
diff --git a/tensorflow/compiler/xla/python/xla_client.pyi b/tensorflow/compiler/xla/python/xla_client.pyi
index 0efe834..c3d8867 100644
--- a/tensorflow/compiler/xla/python/xla_client.pyi
+++ b/tensorflow/compiler/xla/python/xla_client.pyi
@@ -56,6 +56,7 @@
executable: Executable, arguments: Sequence[Sequence[Any]],
backend: Client) -> Sequence[Sequence[numpy.ndarray]]: ...
+def shape_from_pyval(pyval: Any) -> Any: ...
def heap_profile(client: Client) -> bytes:
...
@@ -109,6 +110,10 @@
) -> PaddingConfig:
...
+class PaddingType(enum.Enum):
+ VALID = 1
+ SAME = 2
+
class DotDimensionNumbers:
lhs_contracting_dimensions: List[int]
diff --git a/tensorflow/compiler/xla/python/xla_extension/ops.pyi b/tensorflow/compiler/xla/python/xla_extension/ops.pyi
index f6a4aaa..41474e5 100644
--- a/tensorflow/compiler/xla/python/xla_extension/ops.pyi
+++ b/tensorflow/compiler/xla/python/xla_extension/ops.pyi
@@ -163,7 +163,7 @@
operands: Sequence[XlaOp],
shape: Shape,
opaque: bytes = ...,
- has_side_effects: bool = ...,
+ has_side_effect: bool = ...,
schedule: CustomCallSchedule = ...,
api_version: CustomCallApiVersion = ...) -> XlaOp: ...
def CustomCallWithLayout(
@@ -173,7 +173,7 @@
shape_with_layout: Shape,
operand_shapes_with_layout: Sequence[Shape],
opaque: bytes = ...,
- has_side_effects: bool = ...,
+ has_side_effect: bool = ...,
schedule: CustomCallSchedule = ...,
api_version: CustomCallApiVersion = ...) -> XlaOp: ...
def CustomCallWithAliasing(
@@ -183,7 +183,7 @@
shape_with_layout: Shape,
operand_shapes_with_layout: Sequence[Shape],
opaque: bytes = ...,
- has_side_effects: bool = ...,
+ has_side_effect: bool = ...,
output_operand_aliasing: Sequence[Tuple[ShapeIndex, Tuple[int, ShapeIndex]]] = ...,
literal: _LiteralSlice = ...,
schedule: CustomCallSchedule = ...,
