Defining distributed partials element-by-element doesn't work

distrib_jacobian.py

import openmdao.api as om
import numpy as np

from openmdao.utils.array_utils import evenly_distrib_idxs

N = 3

class DistribCompNoWork(om.ExplicitComponent):

    def initialize(self):
        self.options['distributed'] = True

    def setup(self):
        self.add_input('x', shape=1, src_indices=1)
        rank = self.comm.rank
        sizes, offsets = evenly_distrib_idxs(self.comm.size, N)
        self.add_output('y', shape=sizes[rank])
        self.declare_partials('y', 'x')


    def compute(self, inputs, outputs):
        rank = self.comm.rank
        sizes, offsets = evenly_distrib_idxs(self.comm.size, N)
        outputs['y'] = inputs['x']*np.ones((sizes[rank],))

    def compute_partials(self, inputs, J):
        rank = self.comm.rank
        sizes, offsets = evenly_distrib_idxs(self.comm.size, N)
        # Define jacobian element by element with variable size array
        J['y','x'] = np.ones((sizes[rank],))

class DistribCompNoWork2(om.ExplicitComponent):

    def initialize(self):
        self.options['distributed'] = True

    def setup(self):
        self.add_input('x', shape=1, src_indices=1)
        rank = self.comm.rank
        sizes, offsets = evenly_distrib_idxs(self.comm.size, N)
        self.add_output('y', shape=sizes[rank])
        self.declare_partials('y', 'x')


    def compute(self, inputs, outputs):
        rank = self.comm.rank
        sizes, offsets = evenly_distrib_idxs(self.comm.size, N)
        outputs['y'] = inputs['x']*np.ones((sizes[rank],))

    def compute_partials(self, inputs, J):
        rank = self.comm.rank
        sizes, offsets = evenly_distrib_idxs(self.comm.size, N)
        # Define jacobian element by element with an array equal to global size
        J['y','x'] = np.ones((N,))

class DistribCompWorks(om.ExplicitComponent):

    def initialize(self):
        self.options['distributed'] = True

    def setup(self):
        self.add_input('x', shape=1, src_indices=1)
        rank = self.comm.rank
        sizes, offsets = evenly_distrib_idxs(self.comm.size, N)
        self.add_output('y', shape=sizes[rank])
        self.declare_partials('y', 'x')


    def compute(self, inputs, outputs):
        rank = self.comm.rank
        sizes, offsets = evenly_distrib_idxs(self.comm.size, N)
        outputs['y'] = inputs['x']*np.ones((sizes[rank],))

    def compute_partials(self, inputs, J):
        rank = self.comm.rank
        sizes, offsets = evenly_distrib_idxs(self.comm.size, N)
        # Define jacobian with a scalar
        J['y','x'] = 1.0


class GroupThatDoesNotWork(om.Group):
    def setup(self):
        dvs = om.IndepVarComp()
        dvs.add_output('x', val=6.0)
        self.add_subsystem('dvs', dvs, promotes_outputs=['*'])
        self.add_subsystem('distcomp',DistribCompNoWork(), promotes_inputs=['*'])
        self.add_subsystem('execcomp',om.ExecComp('z = y', y=np.zeros((N,)), z=np.zeros((N,))))
        self.connect('distcomp.y', 'execcomp.y')
        self.add_design_var('x', lower=0.0, upper=10.0, scaler=1.0)
        self.add_constraint('execcomp.z', lower=4.2, scaler=1.0)
        self.add_objective('x')

class GroupThatAlsoDoesNotWork(om.Group):
    def setup(self):
        dvs = om.IndepVarComp()
        dvs.add_output('x', val=6.0)
        self.add_subsystem('dvs', dvs, promotes_outputs=['*'])
        self.add_subsystem('distcomp',DistribCompNoWork2(), promotes_inputs=['*'])
        self.add_subsystem('execcomp',om.ExecComp('z = y', y=np.zeros((N,)), z=np.zeros((N,))))
        self.connect('distcomp.y', 'execcomp.y')
        self.add_design_var('x', lower=0.0, upper=10.0, scaler=1.0)
        self.add_constraint('execcomp.z', lower=4.2, scaler=1.0)
        self.add_objective('x')

class GroupThatDoesWork(om.Group):
    def setup(self):
        dvs = om.IndepVarComp()
        dvs.add_output('x', val=6.0)
        self.add_subsystem('dvs', dvs, promotes_outputs=['*'])
        self.add_subsystem('distcomp',DistribCompWorks(), promotes_inputs=['*'])
        self.add_subsystem('execcomp',om.ExecComp('z = y', y=np.zeros((N,)), z=np.zeros((N,))))
        self.connect('distcomp.y', 'execcomp.y')
        self.add_design_var('x', lower=0.0, upper=10.0, scaler=1.0)
        self.add_constraint('execcomp.z', lower=4.2, scaler=1.0)
        self.add_objective('x')

def mainfunc():
    prob = om.Problem()
    prob.model = GroupThatDoesNotWork()
    # prob.model = GroupThatAlsoDoesNotWork
    # prob.model = GroupThatDoesWork()
    prob.driver = om.pyOptSparseDriver()
    prob.driver.options['optimizer'] = "SLSQP"

    prob.setup(mode='fwd')
    prob['dvs.x'] = 7.0
    prob.run_model()

    prob.check_partials(includes='exec*',compact_print=True)
    # prob.check_totals(compact_print=True)
    prob.run_driver()
    a = prob['execcomp.z']
    print(a)

if __name__ == "__main__":
    mainfunc()