I don’t see any obvious errors.
However, what is your use-case for ident_zeros? There might be better ways of performing the task you have in mind than modifying the PETSc matrix post assembly.
For instance: Modify matrix diagonal -- dolfinx version for A.ident_zeros() - #3 by dokken
illustrates this.
You might also be able to use dolfinx.mesh.create_submesh to restrict the function space instead.