Use "x is y" rather than "id(x) == id(y)".

The latter involves at least two object constructions.

PiperOrigin-RevId: 736878098

Author: hawkinsp

jax/_src/interpreters/mlir.py

@@ -1870,7 +1870,7 @@ def __eq__(self, other):
     if self.value.aval != other.value.aval:
       return False
     if self.data is None:
-      return id(self) == id(other)
+      return self is other
     return self.data == other.data
 
 

jax/_src/mesh.py

@@ -286,7 +286,7 @@ def __reduce__(self):
   def __eq__(self, other):
     # This is a performance optimization. Comparing thousands of devices
     # can be expensive.
-    if id(self) == id(other):
+    if self is other:
       return True
     if not isinstance(other, Mesh):
       return False
@@ -454,7 +454,7 @@ def __hash__(self):
     return self._hash
 
   def __eq__(self, other):
-    if id(self) == id(other):
+    if self is other:
       return True
     if not isinstance(other, AbstractMesh):
       return False

jax/_src/numpy/scalar_types.py

@@ -40,7 +40,7 @@ def __hash__(self) -> int:
     return hash(self.dtype.type)
 
   def __eq__(self, other: Any) -> bool:
-    return id(self) == id(other) or self.dtype.type == other
+    return self is other or self.dtype.type == other
 
   def __ne__(self, other: Any) -> bool:
     return not (self == other)

jax/_src/op_shardings.py

@@ -55,7 +55,7 @@ def is_op_sharding_replicated(op: xc.OpSharding | xc.HloSharding) -> bool:
 
 def are_op_shardings_equal(op1: xc.OpSharding | xc.HloSharding,
                            op2: xc.OpSharding | xc.HloSharding) -> bool:
-  if id(op1) == id(op2):
+  if op1 is op2:
     return True
   if is_op_sharding_replicated(op1) and is_op_sharding_replicated(op2):
     return True