`zeta` doesn't support gradient computation
Its really a Bug?
I think its functionnal issue.  Hi @trickiwoo, 
Thanks for reporting the issue!
I was able to replicate the issue in Colab using TF v2.11 and tf-nightly(2.13.0-dev20230319). Please find the gists  [here(2.11)](https://colab.sandbox.google.com/gist/synandi/057fe7f51c5b1d54d3245f9f46cb4f84/60041_2-11.ipynb) and [here(tf-nightly)](https://colab.sandbox.google.com/gist/synandi/3a1c3095cf86462a482946c27f2951e5/60041.ipynb). It seems like we have to dig more into this issue, we will update soon here. Thank you! Hi @trickiwoo, Apologies for the delay. 

It's an intended behaviour.  `tape.jacobian(t,x)` is returning `None` because of floating point errors. 
As per the formula, 
![image](https://user-images.githubusercontent.com/98147397/227233026-24a76f6d-5c77-4cd7-98c5-cc6e8ea6a2e0.png)

The Hurwitz zeta function involves summing up a series of terms, when you evaluate zeta(x, y) with x=2 and y=2, the function becomes very small for large values of n, and computing the sum numerically may lead to numerical instability or precision errors in the computations. Thank you!

 
 Thanks for your time in investigating the issue! @synandi 

I don't believe that summing up a series of small values can cause numerical instability, as the forward pass appears to work correctly. The only difference between zeta and its gradient is the presence of an additional coefficient in the gradient.

And also, as this [page](https://www.tensorflow.org/api_docs/python/tf/raw_ops) implies, `zeta` supports gradient computation. If it is an intended behavior, it will be appreciated if to clarify it on this page. @trickiwoo, `zeta` supports gradient computation. Please consider this code, where I tried to compute the derivative of `t` with respect to `y`. 
```
import tensorflow as tf
import numpy as np

x = tf.Variable(tf.constant([2], dtype=tf.float32))
y = tf.Variable(tf.constant([2], dtype=tf.float32))

def zeta(x, y):
    t = tf.math.zeta(x, y)
    return t

t = zeta(x, y)
print(t)
with tf.GradientTape() as tape:
    tape.watch(x)
    tape.watch(y)
    t = zeta(x, y)

gradient = tape.jacobian(t, [x,y])
print(gradient)
```

Output is as follows:
```
tf.Tensor([0.6449341], shape=(1,), dtype=float32)
[None, <tf.Tensor: shape=(1, 1), dtype=float32, numpy=array([[-0.40411383]], dtype=float32)>]
```

The first element of the output is `None` because the derivative of `t` with respect to `x` is `0` and `-0.40411383` is the partial derivative of `t` with respect to `y`. Thank you! Hi @synandi , thanks for looking into the issue!

Zeta is differentiable with respect to x at `x=2` and the derivate of t with respect to x is `-0.937548` according to Mathematica results:
![image](https://user-images.githubusercontent.com/121965696/227678352-65468c2a-3528-4d62-b950-de619b4aea2f.png)
 @trickiwoo It's just not implemented for x.  There's a TODO [here](https://github.com/tensorflow/tensorflow/blob/e060e74b0f47c1832179554486ee211185588e89/tensorflow/python/ops/math_grad.py#L1125).

The derivatives w.r.t. `x` are quite complex.  See [the formulae](https://dlmf.nist.gov/25.11#E18).  Is this something you actually need?  This is probably something that would need to be done by an external contributor. i would like to work on this issue can you please assign me this? Calculating the Jacobian and the differential (gradient) of the Riemann zeta function with mpmath is not directly supported as TensorFlow's automatic differentiation. However, we can numerically approximate the derivatives using finite differences. so i think this should address the issue

**import mpmath
def zeta(s):
    return mpmath.zeta(s)
mpmath.mp.dps = 50  # You can adjust the precision as needed
s = mpmath.mpc(2, 3)  # Replace with your desired complex value of s
dx = 1e-8  # Small perturbation for the real part of s
dy = 1e-8j  # Small perturbation for the imaginary part of s
df_ds_real = (zeta(s + dx) - zeta(s)) / dx
df_ds_imag = (zeta(s + dy) - zeta(s)) / dy

jacobian = mpmath.matrix([[df_ds_real], [df_ds_imag]])
print("Jacobian (gradient):")
print(jacobian)**
 @amishhaa , Thanks for your time and effort on working on this issue.
Could you please create a PR with the necessary changes to address the issue. > Calculating the Jacobian and the differential (gradient) of the Riemann zeta function with mpmath is not directly supported as TensorFlow's automatic differentiation. However, we can numerically approximate the derivatives using finite differences. so i think this should address the issue

You're probably going to want those perturbations to be dependent on the magnitude of the input values. You are correct. When calculating numerical derivatives using finite differences, it is essential to choose appropriate perturbation values based on the magnitude of the input values. 
To address this concern, you can scale the perturbations relative to the magnitude of the input values.


import mpmath
def zeta(s):
    return mpmath.zeta(s)
mpmath.mp.dps = 50  # You can adjust the precision as needed
s = mpmath.mpc(2, 3)  # Replace with your desired complex value of s
perturbation_scale = 1e-6  # Adjust the scale as needed
dx = perturbation_scale * abs(s.real)
dy = perturbation_scale * abs(s.imag)
df_ds_real = (zeta(s + dx) - zeta(s - dx)) / (2 * dx)
df_ds_imag = (zeta(s + dy) - zeta(s - dy)) / (2 * dy)

jacobian = mpmath.matrix([[df_ds_real], [df_ds_imag]])
print("Jacobian (gradient):")
print(jacobian)
 @synandi @cantonios @sachinprasadhs 

Hi, I'm new to this repo. I read discussion, my approach: since zeta function is differentiable w.r.t x, we can define a custom gradient like tf.custom_gradient decorator to compute gradients correctly.

Would like to collaborate with any already working on this OR get guidance from experienced contributors in the repo. Thanks! @yokeshwaran1 do you want to collaborate with me on this issue ?
 > @yokeshwaran1 do you want to collaborate with me on this issue ?

Hi @vulkomilev , apologize for the delay in my response; I missed seeing ur message earlier. Yeah, I'm interested in collaborating on this issue Hi @cantonios ,

I'm relatively new to the world of open source contributions, and I'm eager to get involved.If you see any room for improvement or encounter any issues, please let me know.

Regarding the problem itself, could we consider addressing it using the following approach?

```python
@ops.RegisterGradient("Zeta")
def _ZetaGrad(op, grad):
    """Returns the gradient of zeta(x, q) with respect to x and q."""
    x = op.inputs[0]
    q = op.inputs[1]

    # Broadcasting gradients
    sx = array_ops.shape(x)
    sq = array_ops.shape(q)
    unused_rx, rq = gen_array_ops.broadcast_gradient_args(sx, sq)

    # Evaluating gradient
    with ops.control_dependencies([grad]):
        x_conj = math_ops.conj(x)
        q_conj = math_ops.conj(q)
        partial_q = -x_conj * math_ops.zeta(x + 1, q)  # Derivative with respect to q
        **partial_x = -tf.reduce_sum(tf.math.log(q + tf.range(tf.shape(q), dtype=q.dtype)) / (q + tf.range(tf.shape(q), dtype=q.dtype))**x_conj)  # Derivative with respect to x

        return (partial_x,  # Derivative with respect to x
                array_ops.reshape(math_ops.reduce_sum(partial_q * grad, rq), sq))**
 This issue is stale because it has been open for 180 days with no activity. It will be closed if no further activity occurs. Thank you.
<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

source

### Tensorflow Version

2.11

### Custom Code

Yes

### OS Platform and Distribution

_No response_

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

```shell
`zeta` doesn't support gradient computation. The Hurwitz zeta function is differentiable with respect to x.
The partial derivative with respect to x can be expressed as:

dζ(x, q) / dx = - ∑(n=0 to ∞) (n + q)^(-x) * ln(n + q)
```


### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import numpy as np

x = tf.Variable(tf.constant([2], dtype=tf.float32))
y = tf.Variable(tf.constant([2], dtype=tf.float32))

def zeta(x, y):
    t = tf.math.zeta(x, y)
    return t

t = zeta(x, y)
print(t)
with tf.GradientTape() as tape:
    tape.watch(x)
    t = zeta(x, y)

gradient = tape.jacobian(t, x)
print(gradient)
```


### Relevant log output

```shell
tf.Tensor([0.6449342], shape=(1,), dtype=float32)
None
```
</details>