Errors related to `name` argument in tfp.sts package

Hi, thank you for providing a great product. I'm trying to build a model with tfp.sts package but I'm facing issues around the name argument.

`TypeError: tensorflow_probability.python.sts.components.autoregressive.AutoregressiveStateSpaceModel() got multiple values for keyword argument 'name'`

Script

import tensorflow_probability as tfp
tfd = tfp.distributions

if __name__ == "__main__":
    components = [tfp.sts.Autoregressive(order=1, name="ar")]
    components_params = {"ar/_coefficients": [1.], "ar/_level_scale": [1.]}

    s = tfp.sts.Sum(components, name="sum")

    ssm = s.make_state_space_model(
        num_timesteps=100,
        param_vals={
            "observation_noise_scale": 1e-1,
            **components_params,
        },
        name="ssm"
    )
    print(f"{ssm.name} is built successfully")

Error

Traceback (most recent call last):
  File "/Users/hiroki/dev/tensorflow-probability-issue/example.py", line 21, in <module>
    ssm = s.make_state_space_model(
  File "/Users/hiroki/Library/Caches/pypoetry/virtualenvs/tensorflow-probability-issue-_b8RfAk--py3.10/lib/python3.10/site-packages/tensorflow_probability/python/sts/structural_time_series.py", line 244, in make_state_space_model
    return self._make_state_space_model(
  File "/Users/hiroki/Library/Caches/pypoetry/virtualenvs/tensorflow-probability-issue-_b8RfAk--py3.10/lib/python3.10/site-packages/tensorflow_probability/python/sts/components/sum.py", line 555, in _make_state_space_model
    component_ssms = self.make_component_state_space_models(
  File "/Users/hiroki/Library/Caches/pypoetry/virtualenvs/tensorflow-probability-issue-_b8RfAk--py3.10/lib/python3.10/site-packages/tensorflow_probability/python/sts/components/sum.py", line 537, in make_component_state_space_models
    component.make_state_space_model(
  File "/Users/hiroki/Library/Caches/pypoetry/virtualenvs/tensorflow-probability-issue-_b8RfAk--py3.10/lib/python3.10/site-packages/tensorflow_probability/python/sts/structural_time_series.py", line 244, in make_state_space_model
    return self._make_state_space_model(
  File "/Users/hiroki/Library/Caches/pypoetry/virtualenvs/tensorflow-probability-issue-_b8RfAk--py3.10/lib/python3.10/site-packages/tensorflow_probability/python/sts/components/autoregressive.py", line 415, in _make_state_space_model
    return AutoregressiveStateSpaceModel(
TypeError: tensorflow_probability.python.sts.components.autoregressive.AutoregressiveStateSpaceModel() got multiple values for keyword argument 'name'

`ValueError: Field names must be valid identifiers: arima/`

Script

import tensorflow_probability as tfp
tfd = tfp.distributions

if __name__ == "__main__":
    model = tfp.sts.AutoregressiveIntegratedMovingAverage(0, 0, 1, name="arima")
    @tfd.JointDistributionCoroutine
    def joint_distribution():
        yield tfd.JointDistribution.Root(model.make_state_space_model(1, [1.]))
    print(joint_distribution.log_prob([[[0.]]]))

Error

WARNING:tensorflow:From /Users/hiroki/Library/Caches/pypoetry/virtualenvs/tensorflow-probability-issue-_b8RfAk--py3.10/lib/python3.10/site-packages/tensorflow_probability/python/distributions/distribution.py:342: calling MultivariateNormalDiag.__init__ (from tensorflow_probability.python.distributions.mvn_diag) with scale_identity_multiplier is deprecated and will be removed after 2020-01-01.
Instructions for updating:
`scale_identity_multiplier` is deprecated; please combine it into `scale_diag` directly instead.
Traceback (most recent call last):
  File "/Users/hiroki/dev/tensorflow-probability-issue/example.py", line 36, in <module>
    print(joint_distribution.log_prob([[[0.]]]))
  File "/Users/hiroki/Library/Caches/pypoetry/virtualenvs/tensorflow-probability-issue-_b8RfAk--py3.10/lib/python3.10/site-packages/tensorflow_probability/python/distributions/joint_distribution.py", line 897, in log_prob
    return self._call_log_prob(self._resolve_value(*args, **kwargs), name=name)
  File "/Users/hiroki/Library/Caches/pypoetry/virtualenvs/tensorflow-probability-issue-_b8RfAk--py3.10/lib/python3.10/site-packages/tensorflow_probability/python/distributions/joint_distribution.py", line 808, in _resolve_value
    dtype=self.dtype,
  File "/Users/hiroki/Library/Caches/pypoetry/virtualenvs/tensorflow-probability-issue-_b8RfAk--py3.10/lib/python3.10/site-packages/tensorflow_probability/python/distributions/joint_distribution.py", line 388, in dtype
    return self._model_unflatten(
  File "/Users/hiroki/Library/Caches/pypoetry/virtualenvs/tensorflow-probability-issue-_b8RfAk--py3.10/lib/python3.10/site-packages/tensorflow_probability/python/distributions/joint_distribution_coroutine.py", line 361, in _model_unflatten
    return structural_tuple.structtuple(self._flat_resolve_names())(*xs)
  File "/Users/hiroki/Library/Caches/pypoetry/virtualenvs/tensorflow-probability-issue-_b8RfAk--py3.10/lib/python3.10/site-packages/tensorflow_probability/python/internal/structural_tuple.py", line 73, in structtuple
    _validate_field_names(field_names)
  File "/Users/hiroki/Library/Caches/pypoetry/virtualenvs/tensorflow-probability-issue-_b8RfAk--py3.10/lib/python3.10/site-packages/tensorflow_probability/python/internal/structural_tuple.py", line 48, in _validate_field_names
    raise ValueError('Field names must be valid identifiers: {}'.format(name))
ValueError: Field names must be valid identifiers: arima/

Minimal reproducible example

I created a project to reproduce these errors. After checking out, please install the dependencies:

For M1 mac: poetry install --without tf --with tf_apple_silicon
For other env: poetry install
- Please note that I'm using an M1 mac so this version isn't tested well.

Then, you can reproduce the above errors on each branch:

multiple-name-args branch
- TypeError: tensorflow_probability.python.sts.components.autoregressive.AutoregressiveStateSpaceModel() got multiple values for keyword argument 'name'
- poetry run python example.py => Reproduce the above error: tfp.sts.Autoregressive causes the error
- poetry run python example.py successful => Normal case: tfp.sts.LocalLevel works without any error
invalid-name branch
- ValueError: Field names must be valid identifiers: arima/
- poetry run python example.py => Reproduce the above error: The combination of tfp.sts.AutoregressiveIntegratedMovingAverage and tfd.JointDistributionCoroutine causes the error
- poetry run python example.py locallevel => Normal case: tf.sts.LocalLevel&tfd.JointDistributionCoroutine works without any error
- poetry run python example.py sequential => Normal case: tf.sts.AutoregressiveIntegratedMovingAverage&tfd.JointDistributionSequential works without any error
- poetry run python example.py named => Normal case: tf.sts.AutoregressiveIntegratedMovingAverage&tfd.JointDistributionSequential works without any error

Does someone have any thoughts on how to resolve these issues?

As far as I investigated, both are about the name argument.

For the multiple values for keyword argument 'name' issue, in the _make_state_space_model method,

tfp.sts.Autoregressive passes name=self.name to a state space model class
- https://github.com/tensorflow/probability/blob/dfcf94c162221df135b7ba90a20a8948d41daef9/tensorflow_probability/python/sts/components/autoregressive.py#L418
But, tfp.sts.LocalLevel doesn't do it
- https://github.com/tensorflow/probability/blob/dfcf94c162221df135b7ba90a20a8948d41daef9/tensorflow_probability/python/sts/components/local_level.py#L346-L349

For the tfp.sts.Autoregressive case, if the name key is included in linear_gaussian_ssm_kwargs, the error should occur.

Also, for the ValueError: Field names must be valid identifiers case,

tfp.sts.AutoregressiveIntegratedMovingAverage passes name=self.name to a state space model
- https://github.com/tensorflow/probability/blob/dfcf94c162221df135b7ba90a20a8948d41daef9/tensorflow_probability/python/sts/components/autoregressive_integrated_moving_average.py#L382
And, it is captured as parameters in tfp.sts.AutoregressiveMovingAverageStateSpaceModel
- https://github.com/tensorflow/probability/blob/c0908aba1ac6cbc42ab95e4d43434ab238b0687d/tensorflow_probability/python/sts/components/autoregressive_moving_average.py#L156
Then, the name is used in tfd.JointDistribution
- https://github.com/tensorflow/probability/blob/c0908aba1ac6cbc42ab95e4d43434ab238b0687d/tensorflow_probability/python/distributions/joint_distribution.py#L1164
The name is not valid so the error occurs
- https://github.com/tensorflow/probability/blob/c0908aba1ac6cbc42ab95e4d43434ab238b0687d/tensorflow_probability/python/internal/structural_tuple.py#L39-L56

I wonder how these issues can be resolved. For example, classes under tfp.sts have / suffix in its name:

>>> tfp.sts.AutoregressiveIntegratedMovingAverage(0, 0, 0).name
'ARIMA/'
>>> tfp.sts.AutoregressiveIntegratedMovingAverage(0, 0, 0, name="custom").name
'custom/'
>>> tfp.sts.LocalLevel().name
'LocalLevel/'
>>> tfp.sts.LocalLevel(name="ll").name
'll/'

Also, when making an SSM from tfp.sts.AutoregressiveIntegratedMovingAverage, all the nested models look to have the same name. https://github.com/tensorflow/probability/blob/c0908aba1ac6cbc42ab95e4d43434ab238b0687d/tensorflow_probability/python/sts/components/autoregressive_integrated_moving_average.py#L378-L391

I suspect there is something around these places, but I don't have enough knowledge about the past design decision to consider the right solution. Does anyone have ideas about these points?

tensorflow / probability